Add the Carbon-A annotation pipeline and rewrite the model card around it

INFERENCE.md

@@ -1,86 +0,0 @@
-# Carbon-A inference guide
-
-Examples, output formats, and functional checks for Carbon-A. For model facts and a quick start, see the [model card](README.md).
-
-## Run inference on a GenBank file
-
-Use `infer_genbank.py` for a local eukaryotic GenBank flat file (`.gb`, `.gbff`, or `.gz`). It reads the DNA sequence; existing feature annotations are not used as model inputs. Install the repository's `requirements.txt` and authenticate with `hf auth login`, then run:
-
-```bash
-python infer_genbank.py genome.gbff.gz --out-dir genome-predictions
-```
-
-For a FASTA file downloaded from GenBank:
-
-```bash
-python infer_genbank.py genome.fna.gz --format fasta --out-dir genome-predictions
-```
-
-Both commands process every sequence record in the file. Use `--revision <model-commit>` to pin the model, or `--threshold 0.7` to change the binary-call threshold (default 0.5).
-
-The script uses 98,304-base windows with a 49,152-base stride and averages CDS probabilities where windows overlap. The last window may be shorter. No BOS/EOS tokens are added. Short records are supported, but the training data used full-length 98,304-base windows. This example uses overlap averaging; it does not reproduce any separate production genome-annotation pipeline.
-
-Outputs are one compressed NumPy file per record, named with its accession and record name, plus `summary.json`. Each file contains forward/reverse CDS probabilities and calls, a combined `cds_label_any` track, a `valid` mask, and the number of contributing windows per base (`window_count`). Array position `i` corresponds to record base `i + 1` on **both** strands; the reverse track is not reversed. These are coding-occupancy predictions, not GFF gene or transcript models.
-
-Any 6-mer containing an ambiguous nucleotide is excluded from the output. An incomplete final 6-mer is padded with `N` and also excluded. Unscored positions have probability `NaN`, call `-1`, and `valid=False`. Record sequences must be present in the file; annotation-only records with undefined sequence cannot be used. The example retains per-base arrays for one record in memory, so large chromosomes require more host RAM.
-
-```python
-import json
-from pathlib import Path
-import numpy as np
-
-out = Path("genome-predictions")
-record = json.loads((out / "summary.json").read_text())["records"][0]
-with np.load(out / record["file"]) as prediction:
-    plus = prediction["cds_probability_plus"]
-    minus = prediction["cds_probability_minus"]
-    coding = prediction["cds_label_any"]  # 0, 1, or -1 (unscored)
-    valid = prediction["valid"]
-    print(record["accession"], plus.shape, int(valid.sum()))
-```
-
-## Run inference on packed data
-
-Use Python 3.10+ and a CUDA-enabled PyTorch installation. Install the requirements and authenticate with your Hugging Face account:
-
-```bash
-pip install -r requirements.txt
-hf auth login
-python infer_packed.py --num-rows 1 --out-dir predictions
-```
-
-Run these commands from a downloaded copy of this repository. The script downloads a small evaluation HDF5 file when `--h5` is omitted, loads the model from this repository, and predicts one **full-length** row. Model files are downloaded once and cached. The script uses PyTorch SDPA; a separate FlashAttention package is not required.
-
-For a local training or evaluation shard:
-
-```bash
-python infer_packed.py --h5 /path/to/shard.h5 --start-row 0 --num-rows 2 --batch-size 1 --out-dir predictions-local
-```
-
-For repeatable runs, pass `--revision <model-commit>`. Default precision is BF16 on supported GPUs, otherwise FP16; CPU uses FP32 and requires `--device cpu`. Full-length inference on CPU is slow.
-
-### Output
-
-Each row produces an `.npz` file with arrays of 98,304 positions:
-
-- `cds_probability_plus`, `cds_probability_minus`: CDS probabilities.
-- `cds_label_plus`, `cds_label_minus`: calls at `--threshold` (default 0.5).
-- `cds_label_any`: 1 if either strand passes the threshold.
-- `valid_plus`, `valid_minus`: positions eligible for label comparison.
-- `reference_plus`, `reference_minus`: packed RefSeq targets, converted with `label != 0`.
-
-Predicted calls are `-1` at excluded special-token positions; references use `-100`. Probabilities are retained at these positions, so use the validity masks when comparing predictions with annotations. A JSON sidecar records the model revision, input row indices, precision, threshold, and output paths. Coordinates are relative to the packed window; the HDF5 rows do not supply genomic accession/start metadata.
-
-### Loading details
-
-`load_data.py` reads the already-tokenized inputs without adding BOS/EOS tokens. Raw signed labels are converted to binary with `!= 0`, concatenated forward then reverse, and masked at special token positions. Its token IDs match the model tokenizer. To tokenize raw DNA separately, use uppercase sequence and `add_special_tokens=False`; this script itself accepts packed HDF5 only.
-
-The model requires `trust_remote_code=True` to load its custom architecture. It supports Transformers 4.56 and later, including 5.x. Apache-2.0 notices in the model code are retained.
-
-## Inference check
-
-On 2026-10-02, the published model and script at commit `d5a62f1bbf483426314525fcff8695847f9fcdc4` were downloaded from the Hub and run on rows 0–1 of `test/fungi1/GCF_000002945.2.h5`, fetched from the data bucket by the script. The run passed on one NVIDIA H100 80 GB using BF16, PyTorch 2.7.1+cu128, Transformers 4.56.0, and batch size 2. Each row covered all 98,304 bases, with finite two-strand predictions and checked threshold/validity-mask outputs. This is a functional inference check, not an accuracy benchmark. See `inference-check.json` for the check summary.
-
-### GenBank file inference check
-
-On 2026-10-02, `infer_genbank.py` at commit `d429145e77373827ac93cc03ff04dcbdfab5723f` passed an end-to-end test on an NVIDIA H100 using BF16. The input included the complete 230,218-base yeast chromosome I record `NC_001133.9` (four overlapping windows), plus synthetic short and ambiguous-sequence edge cases. Gzipped GenBank and FASTA inputs produced identical predictions. Saved probabilities matched an independent window-by-window calculation and overlap average; coordinates, thresholded calls, invalid-position masks, multiple records, and overwrite refusal were checked. No script changes were required. See `genbank-inference-check.json`. This check establishes that the example runs correctly; it is not an accuracy benchmark.

README.md

@@ -1,6 +1,6 @@
 ---
 library_name: transformers
-license: mit
+license: apache-2.0
 language:
   - dna
 pipeline_tag: token-classification
@@ -17,134 +17,82 @@ tags:
 
 A DNA annotation model from the **Carbon** family.
 
-**Carbon-A predicts protein-coding sequence (CDS) at single-base resolution in eukaryotic genomes.** It produces separate forward- and reverse-strand probabilities, preserving coding regions that overlap on opposite strands.
+**Carbon-A predicts protein-coding sequence (CDS) at single-base resolution in eukaryotic genomes.** It produces separate forward- and reverse-strand probabilities, preserving coding regions that overlap on opposite strands. The [annotation pipeline](annotation_pipeline/README.md) in this repository turns these probabilities into gene models with a confidence score per gene.
 
 ## Facts
 
 - **1.2B-parameter DNA annotation model** with two strand-specific classification heads.
 - **Tokenizer:** non-overlapping 6-mers; each DNA token represents six bases.
-- **Sequence length:** 16,384 tokens (**98,304 bp**).
-- **Inputs:** genomic DNA. The examples support GenBank, FASTA, and packed HDF5 files.
-- **Outputs:** per-base probabilities for non-coding and CDS on each strand; binary calls use a configurable threshold, default **0.5**.
-- **Training data:** RefSeq GCF assemblies from mammals, other vertebrates, invertebrates, plants, fungi, and protozoa.
-- **Transformers interface:** `AutoModelForTokenClassification` with `trust_remote_code=True`; Transformers 4.56 or later, including 5.x.
-- **Inference:** PyTorch SDPA, with BF16 on supported GPUs. A separate FlashAttention installation is not required.
+- **Sequence length:** 16,384 tokens (**98,304 bp**). The annotation pipeline reads longer records in windows that overlap by 1,024 tokens (6,144 bp).
+- **Inputs:** genomic DNA. The annotation pipeline reads FASTA and GenBank files, optionally gzip-compressed.
+- **Outputs:** per-base probabilities for non-coding and CDS on each strand. The annotation pipeline turns them into gene models with a confidence score per gene.
+- **Training data:** 2,055 RefSeq GCF assemblies (2,047 species) from mammals, other vertebrates, invertebrates, plants, fungi, and protozoa.
+- **Transformers interface:** `AutoModelForTokenClassification` with `trust_remote_code=True`; Transformers 4.56 or later, including 5.x. FlashAttention-2 needs Transformers 4.x.
+- **Inference:** BF16 with FlashAttention-2, the setting of all reported results. PyTorch SDPA (`--attn-implementation sdpa`) needs no FlashAttention installation, but its probabilities are not bit-identical.
 
 ## How to use
 
-Use a Python 3.10+ environment with CUDA-enabled PyTorch. Authenticate with your Hugging Face account, download the example scripts, and install their requirements:
+### Install
+
+The annotation pipeline runs on Linux. It needs Python 3.10 or later, a CUDA build of PyTorch, and a C++ compiler with OpenMP. If you do not have a suitable compiler, we recommend GCC, for example on Debian or Ubuntu:
 
 ```bash
-pip install -U "huggingface_hub>=0.34"
-hf auth login
-hf download HuggingFaceBio/Carbon-A-1.2B --exclude model.safetensors --local-dir Carbon-A-1.2B
-cd Carbon-A-1.2B
-pip install -r requirements.txt
+sudo apt install build-essential
 ```
 
-The scripts load the weights through `from_pretrained`, which downloads them (4.6 GB) to the Hugging Face cache on first use, so they are excluded above.
-
-Full-length 98,304-base windows rely on PyTorch's fused SDPA kernels (FlashAttention or memory-efficient attention), which recent NVIDIA GPUs provide. Without them, PyTorch falls back to a kernel that materializes 16,384 × 16,384 attention matrices and needs more than 32 GB of additional GPU memory.
-
-### Use from Python
-
-```python
-import torch
-from transformers import AutoModelForTokenClassification, AutoTokenizer
-
-model_id = "HuggingFaceBio/Carbon-A-1.2B"
-tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModelForTokenClassification.from_pretrained(
-    model_id, trust_remote_code=True, dtype=torch.bfloat16, attn_implementation="sdpa"
-).to("cuda").eval()
+Then download this repository and install its requirements:
 
-sequence = "ATGTCTGCTACACAAAAGCAACACTAA"  # up to 98,304 bases
-sequence = sequence.upper()
-sequence += "N" * (-len(sequence) % 6)  # complete the final 6-mer
-inputs = tokenizer(sequence, add_special_tokens=False, return_tensors="pt").to("cuda")
-with torch.inference_mode():
-    logits = model(**inputs).logits  # (1, 2 * len(sequence), 2)
-cds_plus, cds_minus = logits.float().softmax(-1)[0, :, 1].reshape(2, -1).cpu()
+```bash
+pip install -U huggingface_hub
+hf download HuggingFaceBio/Carbon-A-1.2B --exclude model.safetensors --local-dir Carbon-A-1.2B
+pip install -r Carbon-A-1.2B/requirements.txt
+pip install --no-build-isolation flash-attn
 ```
 
-Logits are per base: all forward-strand positions, then all reverse-strand positions, both in input coordinates. Positions covered by a 6-mer containing `N` or another ambiguous base are not meaningful; the example scripts exclude them. Pass `attention_mask`, which the tokenizer returns, because the model requires it. Uppercase the sequence and use `add_special_tokens=False`: the tokenizer adds a BOS token by default, and lowercase (soft-masked) bases are not grouped into 6-mers, which shifts every later position.
-
-### Annotate a GenBank or FASTA file
+The weights (4.6 GB) are downloaded to the Hugging Face cache on first use.
 
-The included [infer_genbank.py](infer_genbank.py) script reads the DNA sequence from each record and writes per-base predictions:
+### Annotate a genome
 
 ```bash
-# GenBank flat file; gzip compression is supported.
-python infer_genbank.py genome.gbff.gz --out-dir predictions
-
-# FASTA input.
-python infer_genbank.py genome.fna.gz --format fasta --out-dir predictions-fasta
+bash Carbon-A-1.2B/annotation_pipeline/run_annotation.sh --input genome.fna.gz
 ```
 
-Long records are processed in **98,304-base windows with 50% overlap**. Probabilities are averaged where windows overlap. Existing GenBank annotations are not model inputs.
+The input is a FASTA or GenBank file, optionally gzip-compressed. Carbon-A reads it in 98,304-bp windows that overlap by 6,144 bp and averages the overlapping predictions. The pipeline then decodes the gene structures and writes the results to `results/<accession>/`:
 
-Each record produces an `.npz` file named with its accession and record name, plus an entry in `summary.json`. To read the predictions:
+- `prediction.gff`: genes, mRNAs, and CDS. Every gene carries `confidence`, the probability that its structure is exactly right.
+- `prediction.fna` and `prediction.faa`: the CDS and protein sequences.
+- `*.raw_probs.parquet`: the per-base CDS probabilities.
 
-```python
-import json
-from pathlib import Path
-import numpy as np
+All genes are written. The reported results count the genes with a confidence of at least 0.1.
 
-out = Path("predictions")
-record = json.loads((out / "summary.json").read_text())["records"][0]
-with np.load(out / record["file"]) as pred:
-    plus = pred["cds_probability_plus"]
-    minus = pred["cds_probability_minus"]
-    coding = pred["cds_label_any"]  # 1 if either strand passes the threshold
-    valid = pred["valid"]
-```
+The standard genetic code is used by default. For another code, pass `--codon-table`. *Tetrahymena thermophila* (`GCF_000189635.1`) uses translation table 6 automatically. A likely codon-table mismatch is detected automatically: the pipeline adds a warning to `results/WARNINGS.txt` and renames the result directory to `results/<accession>_WARNING`.
 
-Array index `i` corresponds to record base `i + 1` on both strands. Unscored positions have probability `NaN`, call `-1`, and `valid=False`.
+The [pipeline README](annotation_pipeline/README.md) lists all options. It also describes the evaluation against a reference annotation and BUSCO.
 
-### Run on packed training or evaluation data
+### Test on a benchmark genome
 
-The included [infer_packed.py](infer_packed.py) script can download a small evaluation file and predict one full-length row:
+The benchmark genomes and their NCBI annotations are in the [data bucket](https://huggingface.co/buckets/HuggingFaceBio/Carbon-A-training-data) under `eval/`. To test the pipeline on yeast, download its genome and annotation:
 
 ```bash
-python infer_packed.py --num-rows 1 --out-dir predictions-packed
+url=https://huggingface.co/buckets/HuggingFaceBio/Carbon-A-training-data/resolve/eval/seen/GCF_000146045.2
+mkdir -p GCF_000146045.2
+for file in GCF_000146045.2_R64_genomic.fna genomic.gff cds_from_genomic.fna sequence_report.jsonl; do
+    curl -fsSL -o GCF_000146045.2/$file $url/$file
+done
 ```
 
-Or use a local shard:
+Annotate the genome and compare the result with the NCBI annotation:
 
 ```bash
-python infer_packed.py --h5 /path/to/shard.h5 --num-rows 2 --out-dir predictions-local
+bash Carbon-A-1.2B/annotation_pipeline/run_annotation.sh \
+    --input GCF_000146045.2/GCF_000146045.2_R64_genomic.fna
+bash Carbon-A-1.2B/annotation_pipeline/run_evaluation.sh \
+    --prediction-prefix results/GCF_000146045.2/prediction \
+    --genome-dir GCF_000146045.2
 ```
 
-Both scripts accept `--threshold` and `--revision <model-commit>`. See the **[inference guide](INFERENCE.md)** for output fields, tokenization, coordinate handling, and reproducible-run details.
-
-## Training data
-
-Examples were constructed from paired **RefSeq GBFF annotations and FASTA sequences**, using GCF-accessioned assemblies. Records at least 98,304 bp long were eligible. Domain-specific sliding windows increase exposure to less abundant groups, particularly fungi and protozoa.
-
-CDS intervals are unioned across isoforms independently on each strand. Windows without annotated CDS are retained, exposing the model to both coding and non-coding genomic context.
-
-The **[packed-data bucket](https://huggingface.co/buckets/HuggingFaceBio/Carbon-A-training-data)** contains training shards, the original evaluation panel, a PyTorch loader, and reuse instructions. Stored labels are converted to binary with `label != 0`; the loader handles signed values and special-token masking.
-
-## Evaluation
-
-The inference examples have passed functional checks on an **NVIDIA H100**, using BF16, PyTorch 2.7.1, and Transformers 4.56.0:
-
-- Two full-length packed evaluation windows, including strand outputs and thresholded calls.
-- The complete **230,218-base yeast chromosome I** record, processed in four overlapping windows.
-- Identical predictions when the same sequence is supplied as GenBank or FASTA.
-- Independent checks of overlap averaging, coordinates, ambiguous bases, and short records.
-
-These checks verify the inference code. **They do not measure agreement with reference CDS annotations.** Details are recorded in [the packed-data check](inference-check.json) and [the GenBank check](genbank-inference-check.json).
-
-## Limitations
-
-- Predictions represent **CDS occupancy**. Transcript paths, alternative isoforms, and complete gene structures are not resolved.
-- The model is intended for **eukaryotic sequence**. Bacteria and viruses are outside its training scope.
-- The GenBank/FASTA example excludes any 6-mer containing an ambiguous nucleotide, including an incomplete final 6-mer padded with `N`.
-- Shorter records can be processed, but training used full-length 98,304-base windows.
-- The supplied evaluation panel was used during training; it is not a newly held-out test set.
-- Full-length inference on CPU is slow. The file example holds per-base arrays for one record in memory, so large chromosomes need more host RAM.
+`results/GCF_000146045.2/evaluation_metrics.json` then reports a nucleotide F1 of 0.998, an exon F1 of 0.943, and a gene F1 of 0.953.
 
 ## License
 
-MIT. Apache-2.0 notices in the model code are retained.
+Apache 2.0.

annotation_pipeline/README.md

@@ -0,0 +1,258 @@
+# Carbon-A annotation pipeline
+
+This directory annotates protein-coding genes in eukaryotic genomes with
+Carbon-A: it runs the model over a genome and turns its per-nucleotide coding
+probabilities into gene annotations. It also contains reference-based
+evaluation and BUSCO evaluation.
+
+## Pipeline
+
+Annotation runs four stages for each genome:
+
+1. **Genome-scale inference.** Carbon-A assigns every nucleotide a
+   strand-specific coding probability (`src/model_inference.py`).
+2. **Transition network.** A small temporal convolutional network reads each
+   strand's coding-probability track and predicts, at every boundary between
+   adjacent nucleotides, a distribution over remaining in the current segment,
+   entering a gene, leaving a gene, entering an intron, and leaving an intron
+   (`src/transition_network.py`).
+3. **Viterbi algorithm.** The coding probabilities, the transition probabilities,
+   and the gene grammar are combined into the highest-scoring labeling of each
+   strand into intergenic, exon, and intron bases. A gene opens at a start codon
+   and closes at the first in-frame stop codon of the selected genetic code, an
+   intron opens at a donor site and closes at an acceptor site of a permitted
+   splice class (GT-AG, GC-AG, AT-AC), and the reading frame continues across
+   each intron. Genes truncated by a record edge are kept and marked incomplete
+   (`src/viterbi.py`, native Viterbi backends in `src/viterbi_backends/`).
+4. **Confidence network.** For every complete gene, a second network reads the
+   coding probabilities together with the decoded roles of the whole strand,
+   neighbouring genes included, and returns the probability that the gene
+   structure is exactly right (`src/confidence_network.py`).
+
+The trained weights of both networks are in `weights/`. `src/post_processing.py`
+drives the three post-processing stages and writes the outputs; the module
+docstring of `src/viterbi.py` documents the 30 states, the transition
+rules, and where each transition probability enters the Viterbi score.
+
+## Layout
+
+- `run_annotation.sh` runs model inference and post-processing for one genome
+  file.
+- `run_evaluation.sh` compares one annotation with an NCBI reference annotation.
+- `run_busco.sh` runs BUSCO on the predicted proteins of one annotation.
+- `src/` contains the Python implementation and the CPU/GPU Viterbi backend sources.
+- `src/constants.py` centralizes cross-module data contracts and pipeline
+  parameters.
+- `weights/` contains the transition network and confidence network weights.
+- `codon_tables/` contains the standard and accession-specific translation tables.
+
+## Installation
+
+The pipeline supports Linux and Python 3.10 or later. Download this repository
+without the model weights, install a PyTorch build compatible with the local
+CUDA environment, then install the requirements of the repository and
+FlashAttention:
+
+```bash
+pip install -U "huggingface_hub>=0.34"
+hf download HuggingFaceBio/Carbon-A-1.2B --exclude model.safetensors --local-dir Carbon-A-1.2B
+pip install -r Carbon-A-1.2B/requirements.txt
+pip install --no-build-isolation flash-attn
+```
+
+Model inference loads Carbon-A with `from_pretrained`, which downloads the
+weights (4.6 GB) to the Hugging Face cache on first use. It runs in bfloat16
+with FlashAttention-2, the setting of all reported results. The requirements
+keep Transformers below version 5, which does not load the model with
+FlashAttention-2.
+`--attn-implementation sdpa` runs without the `flash-attn` package, but its
+probabilities are not bit-identical with those of FlashAttention-2.
+
+Two optional command-line tools are not installed by `requirements.txt`:
+
+- The NCBI Datasets CLI (`datasets`) downloads genomes and reference
+  annotations from NCBI.
+- The BUSCO CLI (`busco`) is required only by `run_busco.sh`.
+
+The two networks run with PyTorch on the GPU when one is available, otherwise on
+the CPU. The GPU Viterbi backend requires a CUDA toolkit; the CPU backend
+requires a C++17 compiler with OpenMP support. Both are built on first use. Use
+`--viterbi-backend numba` when neither native backend can be built.
+
+## Annotation
+
+The input is a genome FASTA or GenBank file, optionally gzip-compressed. Run the
+scripts with `bash`, because `hf download` does not keep their executable bit.
+For example, for the yeast reference genome from NCBI:
+
+```bash
+datasets download genome accession GCF_000146045.2 --include genome,gff3,cds,seq-report
+unzip ncbi_dataset.zip
+bash Carbon-A-1.2B/annotation_pipeline/run_annotation.sh \
+    --input ncbi_dataset/data/GCF_000146045.2/GCF_000146045.2_R64_genomic.fna
+```
+
+The results are written below the current directory:
+
+```text
+results/<accession>/
+  *.raw_probs.parquet
+  prediction.gff
+  prediction.fna
+  prediction.faa
+```
+
+The accession is `--accession` when given, otherwise the NCBI assembly accession
+that begins the file name (`GCF_`/`GCA_`), otherwise the file name without its
+extensions; `--output-dir` selects another directory. Model inference uses every
+visible GPU unless `--gpu-count` is given. `--model` selects another model, a
+Hub id or a local directory (default: `CARBON_A_MODEL`, otherwise
+`HuggingFaceBio/Carbon-A-1.2B`), and `--revision` a branch, tag, or commit of a
+model on the Hub (default: the latest). `run_annotation.sh --help` lists every
+option.
+
+When the codon-table check described below fails, the post-processing itself
+appends the warning to `results/WARNINGS.txt` (time, accession, message with
+the stop codons used) and renames the directory to
+`results/<accession>_WARNING`, whether it runs through the script or on its
+own, so one file and the directory names show every genome whose table is in
+doubt. Rerunning the accession with the right `--codon-table` (usually with
+`--skip-inference`) continues in the renamed directory, renames it back once
+the check passes, and records that in `results/WARNINGS.txt` too.
+
+Every gene row of `prediction.gff` and every FASTA header carries
+`confidence=<score>`, the confidence network's probability that the gene
+structure is exactly right, or `confidence=NA` for genes truncated by a record
+edge. All genes are written; the gene-level results count the genes whose
+confidence reaches 0.1 (see Reference evaluation).
+
+`run_annotation.sh` does not run reference evaluation or BUSCO. By default,
+post-processing excludes records shorter than 100,000 bp and spliced CDS
+models shorter than 96 bp. Set `--min-record-length 0` or `--min-cds-length 0`
+to disable the corresponding filter. Use `--skip-inference` to rerun only
+post-processing from an existing raw-probability Parquet file.
+
+Post-processing options: `--transition-network FILE` and
+`--confidence-network FILE` select other weight files and `--device DEVICE`
+picks the torch device for the two networks (default: `cuda` when available).
+On CUDA the networks run under bfloat16 autocast, the precision used for all
+reported results; on CPU they run in float32. `--viterbi-backend` selects the
+Viterbi backend. `src/post_processing.py` can also be run directly on a
+raw-probability Parquet file.
+
+`GCF_000189635.1` (Tetrahymena thermophila) automatically uses NCBI translation
+table 6. Other accessions use the standard translation table unless an
+accession-specific table is present in `codon_tables/` or supplied with
+`--codon-table`. After scoring, the post-processing checks the codon table with
+the confidences themselves: a table that does not match the genome splits genes
+at codons the genome reads through or fuses genes at stop codons it lacks, and
+the confidence network gives such genes almost no confidence. When at least 100
+complete genes were scored and fewer than half of them reach confidence 0.1,
+it warns that the table probably does not match. With the shipped network, 17
+genomes decoded with their own table keep 67% (hagfish) to 98% of their
+complete genes at or above 0.1, Tetrahymena decoded with the standard table
+keeps 1.4%, and yeast decoded with translation table 6 keeps 36%. A table that
+only drops a rarely used stop codon escapes the check (Dictyostelium decoded
+without TGA keeps 94%). The warning goes to stderr, to `WARNINGS.txt` in the
+parent of the output directory, and into the output directory's name, as
+described above. The summary line reports the number of complete genes at or
+above 0.1 and the median confidence.
+
+## Reference evaluation
+
+Compare an annotation with the NCBI reference annotation of the same assembly:
+
+```bash
+bash Carbon-A-1.2B/annotation_pipeline/run_evaluation.sh \
+    --prediction-prefix results/GCF_000146045.2/prediction \
+    --genome-dir ncbi_dataset/data/GCF_000146045.2
+```
+
+The reference directory provides `genomic.gff` and `cds_from_genomic.fna`. When
+it also provides `sequence_report.jsonl`, evaluation prefers assembled
+chromosomes, then assembled linkage groups, then chromosome scaffolds; without
+the report every record is evaluated. The metrics are written beside the
+predictions, to `results/<accession>/evaluation_metrics.json`, unless
+`--output-json` is given.
+
+The benchmark genomes are also available in this layout, with their NCBI
+reference annotations, under `eval/` in the
+[data bucket](https://huggingface.co/buckets/HuggingFaceBio/Carbon-A-training-data);
+the model card shows how to test the pipeline on one of them.
+
+Gene and protein metrics count the predicted genes whose confidence reaches
+`--confidence-threshold` (default `0.1`). Genes with `confidence=NA` (genes
+truncated by a record edge) are always counted, and `0` counts every predicted
+gene. Exon, boundary, and nucleotide metrics describe the whole annotation.
+`prediction.*` is never modified. The translation table is
+`codon_tables/<accession>.json` when present, where the accession is the name
+of the reference directory, otherwise the standard table; `--codon-table`
+overrides it.
+
+The reported levels are defined as follows:
+
+- Nucleotide metrics measure CDS base overlap.
+- Exon metrics require exact merged CDS intervals from the selected gene model.
+- Gene metrics require a one-to-one exact CDS sequence match on the same genomic
+  record.
+- Protein metrics retain all exact gene matches, then translate unmatched
+  predicted and reference CDS sequences with the same translation table and add
+  one-to-one exact protein matches on the same record. This allows synonymous
+  CDS differences to match while preserving confirmed CDS matches with partial
+  frame annotation differences. NCBI `protein.faa` is not used.
+
+`confidence_mean` is the mean confidence of the counted genes.
+`confidence_exact_match_auroc` and `confidence_exact_match_auprc` treat each
+counted gene as one sample and its CDS-level exact-match status as the binary
+label. Tied scores are processed at the same cutoff. AUPRC is
+non-interpolated average precision. Both values are `NA` when only one label
+class is present.
+
+## BUSCO evaluation
+
+BUSCO must be installed separately and available as `busco`, or supplied with
+`--busco-executable`. It runs offline, so download the lineage dataset first:
+
+```bash
+busco --download saccharomycetaceae_odb12.2
+bash Carbon-A-1.2B/annotation_pipeline/run_busco.sh \
+    --faa results/GCF_000146045.2/prediction.faa \
+    --lineage-path busco_downloads/lineages/saccharomycetaceae_odb12.2
+```
+
+The lineage directory must keep the dataset's name: the script takes the
+lineage name from it and checks it against BUSCO's summary. BUSCO runs in
+protein mode on the proteins in `prediction.faa`. Its working directory is
+temporary; only this compact output is retained, beside the proteins unless
+`--output-json` is given:
+
+```text
+results/<accession>/busco_metrics.json
+```
+
+The accession recorded with the metrics is `--accession`, otherwise the name of
+the directory containing the proteins without a `_WARNING` suffix. `--workers`
+is passed to BUSCO as its CPU count (default 8). An existing
+`busco_metrics.json` is kept unless `--force` is supplied.
+
+The BUSCO JSON records the lineage, its total number of marker genes, and the
+Complete, Single-copy, Duplicated, Fragmented, and Missing percentages. The
+marker count gives important context: scores from a small lineage dataset are
+more sensitive to the state of each individual marker.
+
+## Results layout
+
+After all stages, each genome uses this layout:
+
+```text
+results/<accession>/
+  *.raw_probs.parquet
+  prediction.gff
+  prediction.fna
+  prediction.faa
+  evaluation_metrics.json
+  busco_metrics.json
+```
+
+Both JSON files record the accession, the generation time, and their values
+under `metrics`.

annotation_pipeline/codon_tables/GCF_000189635.1.json

@@ -0,0 +1,18 @@
+{
+  "AAA": "K", "AAT": "N", "AAC": "N", "AAG": "K",
+  "ATA": "I", "ATT": "I", "ATC": "I", "ATG": "M",
+  "ACA": "T", "ACT": "T", "ACC": "T", "ACG": "T",
+  "AGA": "R", "AGT": "S", "AGC": "S", "AGG": "R",
+  "TAA": "Q", "TAT": "Y", "TAC": "Y", "TAG": "Q",
+  "TTA": "L", "TTT": "F", "TTC": "F", "TTG": "L",
+  "TCA": "S", "TCT": "S", "TCC": "S", "TCG": "S",
+  "TGA": "*", "TGT": "C", "TGC": "C", "TGG": "W",
+  "CAA": "Q", "CAT": "H", "CAC": "H", "CAG": "Q",
+  "CTA": "L", "CTT": "L", "CTC": "L", "CTG": "L",
+  "CCA": "P", "CCT": "P", "CCC": "P", "CCG": "P",
+  "CGA": "R", "CGT": "R", "CGC": "R", "CGG": "R",
+  "GAA": "E", "GAT": "D", "GAC": "D", "GAG": "E",
+  "GTA": "V", "GTT": "V", "GTC": "V", "GTG": "V",
+  "GCA": "A", "GCT": "A", "GCC": "A", "GCG": "A",
+  "GGA": "G", "GGT": "G", "GGC": "G", "GGG": "G"
+}

annotation_pipeline/codon_tables/standard.json

@@ -0,0 +1,18 @@
+{
+  "AAA": "K", "AAT": "N", "AAC": "N", "AAG": "K",
+  "ATA": "I", "ATT": "I", "ATC": "I", "ATG": "M",
+  "ACA": "T", "ACT": "T", "ACC": "T", "ACG": "T",
+  "AGA": "R", "AGT": "S", "AGC": "S", "AGG": "R",
+  "TAA": "*", "TAT": "Y", "TAC": "Y", "TAG": "*",
+  "TTA": "L", "TTT": "F", "TTC": "F", "TTG": "L",
+  "TCA": "S", "TCT": "S", "TCC": "S", "TCG": "S",
+  "TGA": "*", "TGT": "C", "TGC": "C", "TGG": "W",
+  "CAA": "Q", "CAT": "H", "CAC": "H", "CAG": "Q",
+  "CTA": "L", "CTT": "L", "CTC": "L", "CTG": "L",
+  "CCA": "P", "CCT": "P", "CCC": "P", "CCG": "P",
+  "CGA": "R", "CGT": "R", "CGC": "R", "CGG": "R",
+  "GAA": "E", "GAT": "D", "GAC": "D", "GAG": "E",
+  "GTA": "V", "GTT": "V", "GTC": "V", "GTG": "V",
+  "GCA": "A", "GCT": "A", "GCC": "A", "GCG": "A",
+  "GGA": "G", "GGT": "G", "GGC": "G", "GGG": "G"
+}

annotation_pipeline/run_annotation.sh

@@ -0,0 +1,340 @@
+#!/usr/bin/env bash
+# Run model inference and post-processing for one genome file.
+set -euo pipefail
+shopt -s nullglob
+
+ROOT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+# Results go below the current directory, never into the pipeline directory.
+RESULTS_DIR="results"
+CODON_TABLE_DIR="$ROOT_DIR/codon_tables"
+# post_processing.py renames results/<accession> to results/<accession>_WARNING
+# and appends the warning to results/WARNINGS.txt when the codon-table check
+# fails (constants.WARNING_DIRECTORY_SUFFIX).
+WARNING_SUFFIX="_WARNING"
+
+INPUT=""
+ACCESSION=""
+MODEL="${CARBON_A_MODEL:-}"
+REVISION=""
+OUTPUT_DIR=""
+CODON_TABLE=""
+VITERBI_BACKEND="auto"
+
+GPU_COUNT="-1"
+BATCH_SIZE="8"
+DISPATCH_CHUNK_SIZE="4"
+ATTN_IMPLEMENTATION="flash_attention_2"
+SKIP_INFERENCE=0
+
+TRANSITION_NETWORK=""
+CONFIDENCE_NETWORK=""
+DEVICE=""
+MIN_RECORD_LENGTH="100000"
+MIN_CDS_LENGTH="96"
+MAX_BOUNDARY_INTRON_LENGTH="20000"
+CHUNK_MBP="500"
+
+PREDICTION_PREFIX=""
+RAW_PROBS=""
+
+usage() {
+    local status="${1:-2}"
+    cat >&2 <<EOF
+Usage:
+  $(basename "$0") --input FILE [options]
+
+Input and output:
+  --input FILE                  Genome FASTA or GenBank file, optionally gzip-compressed.
+  --accession ACCESSION         Result identifier. Default: the NCBI assembly accession
+                                that starts the file name (GCF_/GCA_), otherwise the
+                                file name without its extension.
+  --model MODEL                 Model id on the Hugging Face Hub or local model directory.
+                                Default: CARBON_A_MODEL, otherwise HuggingFaceBio/Carbon-A-1.2B.
+  --revision REVISION           Model revision on the Hub (branch, tag, or commit);
+                                default: the latest.
+  --output-dir DIR              Default: results/<accession> in the current directory, or
+                                results/<accession>_WARNING when a failed codon-table
+                                check renamed it and only that exists.
+  --codon-table FILE            Default: accession-specific table, then standard.
+
+Inference:
+  --gpu-count N                 GPUs to use; -1 selects all available GPUs (default: -1).
+  --batch-size N                Inference batch size (default: 8).
+  --dispatch-chunk-size N       Per-GPU dispatch chunk size (default: 4).
+  --attn-implementation NAME    flash_attention_2 (default; used for all reported results)
+                                or sdpa (needs no flash-attn package; not bit-identical).
+  --skip-inference              Reuse the newest *.raw_probs.parquet in --output-dir.
+
+Post-processing:
+  --viterbi-backend NAME        Viterbi backend: auto, gpu, cpu, or numba (default: auto).
+  --transition-network FILE     Transition network weights
+                                (default: weights/transition_network.safetensors).
+  --confidence-network FILE     Confidence network weights
+                                (default: weights/confidence_network.safetensors).
+  --device DEVICE               Torch device for the two networks (default: cuda if available).
+  --min-record-length N         Minimum FASTA record length (default: 100000).
+  --min-cds-length N            Minimum spliced CDS length (default: 96).
+  --max-boundary-intron-length N
+                                Maximum boundary intron length (default: 20000).
+  --chunk-mbp N                 Post-processing chunk size in Mbp (default: 500).
+
+  -h, --help                    Show this help message.
+EOF
+    exit "$status"
+}
+
+die() {
+    printf 'Error: %s\n' "$*" >&2
+    exit 1
+}
+
+require_value() {
+    local option="$1"
+    local argument_count="$2"
+    (( argument_count >= 2 )) || die "$option requires a value"
+}
+
+parse_args() {
+    while (($#)); do
+        case "$1" in
+            --input)
+                require_value "$1" "$#"
+                INPUT="$2"
+                shift 2
+                ;;
+            --accession)
+                require_value "$1" "$#"
+                ACCESSION="$2"
+                shift 2
+                ;;
+            --model|--model_name)
+                require_value "$1" "$#"
+                MODEL="$2"
+                shift 2
+                ;;
+            --output-dir|--output_path)
+                require_value "$1" "$#"
+                OUTPUT_DIR="$2"
+                shift 2
+                ;;
+            --codon-table|--codon_table)
+                require_value "$1" "$#"
+                CODON_TABLE="$2"
+                shift 2
+                ;;
+            --viterbi-backend|--viterbi_backend)
+                require_value "$1" "$#"
+                VITERBI_BACKEND="$2"
+                shift 2
+                ;;
+            --transition-network|--transition_network)
+                require_value "$1" "$#"
+                TRANSITION_NETWORK="$2"
+                shift 2
+                ;;
+            --confidence-network|--confidence_network)
+                require_value "$1" "$#"
+                CONFIDENCE_NETWORK="$2"
+                shift 2
+                ;;
+            --device)
+                require_value "$1" "$#"
+                DEVICE="$2"
+                shift 2
+                ;;
+            --min-record-length|--min_record_length)
+                require_value "$1" "$#"
+                MIN_RECORD_LENGTH="$2"
+                shift 2
+                ;;
+            --min-cds-length|--min_cds_length)
+                require_value "$1" "$#"
+                MIN_CDS_LENGTH="$2"
+                shift 2
+                ;;
+            --max-boundary-intron-length|--max_boundary_intron_length)
+                require_value "$1" "$#"
+                MAX_BOUNDARY_INTRON_LENGTH="$2"
+                shift 2
+                ;;
+            --chunk-mbp|--chunk_mbp)
+                require_value "$1" "$#"
+                CHUNK_MBP="$2"
+                shift 2
+                ;;
+            --batch-size|--batch_size)
+                require_value "$1" "$#"
+                BATCH_SIZE="$2"
+                shift 2
+                ;;
+            --dispatch-chunk-size|--dispatch_chunk_size)
+                require_value "$1" "$#"
+                DISPATCH_CHUNK_SIZE="$2"
+                shift 2
+                ;;
+            --gpu-count|--gpu_count)
+                require_value "$1" "$#"
+                GPU_COUNT="$2"
+                shift 2
+                ;;
+            --revision)
+                require_value "$1" "$#"
+                REVISION="$2"
+                shift 2
+                ;;
+            --attn-implementation|--attn_implementation)
+                require_value "$1" "$#"
+                ATTN_IMPLEMENTATION="$2"
+                shift 2
+                ;;
+            --skip-inference)
+                SKIP_INFERENCE=1
+                shift
+                ;;
+            -h|--help)
+                usage 0
+                ;;
+            *)
+                printf 'Unknown argument: %s\n\n' "$1" >&2
+                usage 2
+                ;;
+        esac
+    done
+}
+
+accession_from_input() {
+    local filename
+    filename="$(basename -- "$INPUT")"
+
+    # NCBI genome files start with the assembly accession, as in
+    # GCF_000146045.2_R64_genomic.fna.
+    if [[ "$filename" =~ ^(GC[AF]_[0-9]+\.[0-9]+)_ ]]; then
+        printf '%s\n' "${BASH_REMATCH[1]}"
+        return
+    fi
+
+    filename="${filename%.gz}"
+    case "$filename" in
+        *.fasta) printf '%s\n' "${filename%.fasta}" ;;
+        *.fna)   printf '%s\n' "${filename%.fna}" ;;
+        *.fa)    printf '%s\n' "${filename%.fa}" ;;
+        *.gbff)  printf '%s\n' "${filename%.gbff}" ;;
+        *.gbk)   printf '%s\n' "${filename%.gbk}" ;;
+        *.gb)    printf '%s\n' "${filename%.gb}" ;;
+        *)       printf '%s\n' "$filename" ;;
+    esac
+}
+
+resolve_input() {
+    [[ -f "$INPUT" ]] || die "input file does not exist: $INPUT"
+    [[ -n "$ACCESSION" ]] || ACCESSION="$(accession_from_input)"
+}
+
+resolve_paths() {
+    [[ "$ACCESSION" =~ ^[A-Za-z0-9][A-Za-z0-9._-]*$ ]] || \
+        die "invalid accession: $ACCESSION"
+
+    if [[ -z "$OUTPUT_DIR" ]]; then
+        OUTPUT_DIR="$RESULTS_DIR/$ACCESSION"
+        # A rerun (typically --skip-inference with the right --codon-table)
+        # continues in the directory a failed codon-table check renamed.
+        if [[ ! -d "$OUTPUT_DIR" && -d "$OUTPUT_DIR$WARNING_SUFFIX" ]]; then
+            OUTPUT_DIR="$OUTPUT_DIR$WARNING_SUFFIX"
+        fi
+    fi
+    PREDICTION_PREFIX="$OUTPUT_DIR/prediction"
+
+    if [[ -z "$CODON_TABLE" ]]; then
+        if [[ -f "$CODON_TABLE_DIR/$ACCESSION.json" ]]; then
+            CODON_TABLE="$CODON_TABLE_DIR/$ACCESSION.json"
+        else
+            CODON_TABLE="$CODON_TABLE_DIR/standard.json"
+        fi
+    fi
+
+    [[ -f "$CODON_TABLE" ]] || die "codon table does not exist: $CODON_TABLE"
+
+    mkdir -p -- "$OUTPUT_DIR"
+}
+
+run_inference() {
+    (( SKIP_INFERENCE )) && return
+
+    local inference_args=(
+        --input "$INPUT"
+        --output_path "$OUTPUT_DIR"
+        --gpu_count "$GPU_COUNT"
+        --batch_size "$BATCH_SIZE"
+        --dispatch_chunk_size "$DISPATCH_CHUNK_SIZE"
+        --attn_implementation "$ATTN_IMPLEMENTATION"
+    )
+    [[ -n "$MODEL" ]] && inference_args+=(--model_name "$MODEL")
+    [[ -n "$REVISION" ]] && inference_args+=(--revision "$REVISION")
+
+    echo "[annotation] inference: $ACCESSION"
+    python "$ROOT_DIR/src/model_inference.py" "${inference_args[@]}"
+}
+
+select_raw_probs() {
+    local candidates=("$OUTPUT_DIR"/*.raw_probs.parquet)
+    (( ${#candidates[@]} > 0 )) || \
+        die "no raw-probability Parquet found in $OUTPUT_DIR"
+
+    RAW_PROBS="${candidates[0]}"
+    local candidate
+    for candidate in "${candidates[@]:1}"; do
+        [[ "$candidate" -nt "$RAW_PROBS" ]] && RAW_PROBS="$candidate"
+    done
+
+    if (( ${#candidates[@]} > 1 )); then
+        echo "[annotation] found ${#candidates[@]} probability files; using $(basename -- "$RAW_PROBS")"
+    fi
+}
+
+run_postprocessing() {
+    local postprocess_args=(
+        --input "$RAW_PROBS"
+        --output "$PREDICTION_PREFIX"
+        --viterbi_backend "$VITERBI_BACKEND"
+        --codon_table "$CODON_TABLE"
+        --max_boundary_intron_length "$MAX_BOUNDARY_INTRON_LENGTH"
+        --chunk_mbp "$CHUNK_MBP"
+        --min_record_length "$MIN_RECORD_LENGTH"
+        --min_cds_length "$MIN_CDS_LENGTH"
+    )
+    [[ -n "$TRANSITION_NETWORK" ]] && postprocess_args+=(--transition_network "$TRANSITION_NETWORK")
+    [[ -n "$CONFIDENCE_NETWORK" ]] && postprocess_args+=(--confidence_network "$CONFIDENCE_NETWORK")
+    [[ -n "$DEVICE" ]] && postprocess_args+=(--device "$DEVICE")
+
+    echo "[annotation] post-processing: $ACCESSION"
+    python "$ROOT_DIR/src/post_processing.py" "${postprocess_args[@]}"
+}
+
+# post_processing.py records a failed codon-table check itself: it appends the
+# warning to WARNINGS.txt beside the result directories and renames the output
+# directory with the warning suffix (and back when a later run passes), so the
+# final report has to look the directory up again.
+locate_output_dir() {
+    local clean_dir="${OUTPUT_DIR%"$WARNING_SUFFIX"}"
+    if [[ -d "$clean_dir" ]]; then
+        OUTPUT_DIR="$clean_dir"
+    elif [[ -d "$clean_dir$WARNING_SUFFIX" ]]; then
+        OUTPUT_DIR="$clean_dir$WARNING_SUFFIX"
+    fi
+}
+
+main() {
+    parse_args "$@"
+    [[ -n "$INPUT" ]] || usage 2
+
+    resolve_input
+    resolve_paths
+    run_inference
+    select_raw_probs
+    run_postprocessing
+    locate_output_dir
+
+    echo "[annotation] complete: $OUTPUT_DIR"
+}
+
+main "$@"

annotation_pipeline/run_busco.sh

@@ -0,0 +1,147 @@
+#!/usr/bin/env bash
+# Run BUSCO on the predicted proteins of one annotation and retain only busco_metrics.json.
+set -euo pipefail
+
+ROOT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+# A failed codon-table check renames results/<accession> with this suffix
+# (constants.WARNING_DIRECTORY_SUFFIX); the default accession drops it.
+WARNING_SUFFIX="_WARNING"
+
+BUSCO_EXECUTABLE="${BUSCO_EXECUTABLE:-busco}"
+
+FAA=""
+LINEAGE_PATH=""
+ACCESSION=""
+OUTPUT_JSON=""
+FORCE=0
+WORKERS="8"
+TEMPORARY_DIRECTORY=""
+
+usage() {
+    local status="${1:-2}"
+    cat >&2 <<EOF
+Usage:
+  $(basename "$0") --faa FILE --lineage-path DIR [options]
+
+Input and output:
+  --faa FILE                    Predicted proteins, for example
+                                results/<accession>/prediction.faa.
+  --lineage-path DIR            Local BUSCO lineage dataset, for example
+                                busco_downloads/lineages/saccharomycetes_odb12.
+  --accession ACCESSION         Name recorded in the metrics (default: the
+                                directory containing --faa).
+  --output-json FILE            Default: busco_metrics.json beside --faa.
+  --force                       Replace an existing output JSON.
+
+BUSCO configuration:
+  --busco-executable FILE       Default: BUSCO_EXECUTABLE or busco from PATH.
+  --workers N                   CPU workers passed to BUSCO (default: 8).
+  --temporary-directory DIR     Parent directory for temporary BUSCO files.
+
+  -h, --help                    Show this help message.
+EOF
+    exit "$status"
+}
+
+die() {
+    printf 'Error: %s\n' "$*" >&2
+    exit 1
+}
+
+require_value() {
+    local option="$1"
+    local argument_count="$2"
+    (( argument_count >= 2 )) || die "$option requires a value"
+}
+
+parse_args() {
+    while (($#)); do
+        case "$1" in
+            --faa)
+                require_value "$1" "$#"
+                FAA="$2"
+                shift 2
+                ;;
+            --lineage-path)
+                require_value "$1" "$#"
+                LINEAGE_PATH="$2"
+                shift 2
+                ;;
+            --accession)
+                require_value "$1" "$#"
+                ACCESSION="$2"
+                shift 2
+                ;;
+            --output-json)
+                require_value "$1" "$#"
+                OUTPUT_JSON="$2"
+                shift 2
+                ;;
+            --busco-executable)
+                require_value "$1" "$#"
+                BUSCO_EXECUTABLE="$2"
+                shift 2
+                ;;
+            --workers)
+                require_value "$1" "$#"
+                WORKERS="$2"
+                shift 2
+                ;;
+            --temporary-directory)
+                require_value "$1" "$#"
+                TEMPORARY_DIRECTORY="$2"
+                shift 2
+                ;;
+            --force)
+                FORCE=1
+                shift
+                ;;
+            -h|--help)
+                usage 0
+                ;;
+            *)
+                printf 'Unknown argument: %s\n\n' "$1" >&2
+                usage 2
+                ;;
+        esac
+    done
+}
+
+main() {
+    parse_args "$@"
+    [[ -n "$FAA" && -n "$LINEAGE_PATH" ]] || usage 2
+    [[ "$WORKERS" =~ ^[1-9][0-9]*$ ]] || die "--workers must be a positive integer"
+    [[ -s "$FAA" ]] || die "missing protein FASTA: $FAA"
+    [[ -d "$LINEAGE_PATH" ]] || die "lineage directory does not exist: $LINEAGE_PATH"
+
+    if [[ -z "$ACCESSION" ]]; then
+        ACCESSION="$(basename -- "$(dirname -- "$FAA")")"
+        ACCESSION="${ACCESSION%"$WARNING_SUFFIX"}"
+    fi
+    local lineage_name
+    lineage_name="$(basename -- "$LINEAGE_PATH")"
+    local output_json="${OUTPUT_JSON:-$(dirname -- "$FAA")/busco_metrics.json}"
+
+    if [[ -s "$output_json" ]] && (( ! FORCE )); then
+        echo "[busco] skipped $ACCESSION ($output_json exists; use --force)"
+        return 0
+    fi
+
+    local busco_args=(
+        --accession "$ACCESSION"
+        --faa "$FAA"
+        --lineage "$lineage_name"
+        --lineage-path "$LINEAGE_PATH"
+        --busco-executable "$BUSCO_EXECUTABLE"
+        --workers "$WORKERS"
+        --output-json "$output_json"
+    )
+    if [[ -n "$TEMPORARY_DIRECTORY" ]]; then
+        busco_args+=(--temporary-directory "$TEMPORARY_DIRECTORY")
+    fi
+
+    echo "[busco] running $ACCESSION ($lineage_name)"
+    python "$ROOT_DIR/src/busco_evaluation.py" "${busco_args[@]}"
+}
+
+main "$@"

annotation_pipeline/run_evaluation.sh

@@ -0,0 +1,113 @@
+#!/usr/bin/env bash
+# Write the reference metrics of one annotation to evaluation_metrics.json.
+set -euo pipefail
+
+ROOT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+
+PREDICTION_PREFIX=""
+GENOME_DIR=""
+OUTPUT_JSON=""
+CODON_TABLE=""
+CONFIDENCE_THRESHOLD="0.1"
+
+usage() {
+    local status="${1:-2}"
+    cat >&2 <<EOF
+Usage:
+  $(basename "$0") --prediction-prefix PREFIX --genome-dir DIR [options]
+
+Input and output:
+  --prediction-prefix PREFIX    Prediction files without .gff/.fna suffixes,
+                                for example results/<accession>/prediction.
+  --genome-dir DIR              NCBI genome directory containing the reference
+                                genomic.gff and cds_from_genomic.fna; an optional
+                                sequence_report.jsonl enables nuclear-record selection.
+  --output-json FILE            Default: evaluation_metrics.json beside the predictions.
+  --codon-table FILE            Translation table JSON (default: the
+                                accession-specific table, then standard).
+
+Metrics:
+  --confidence-threshold VALUE  Gene-level confidence threshold (default: 0.1;
+                                0 counts every predicted gene).
+
+  -h, --help                    Show this help message.
+EOF
+    exit "$status"
+}
+
+die() {
+    printf 'Error: %s\n' "$*" >&2
+    exit 1
+}
+
+require_value() {
+    local option="$1"
+    local argument_count="$2"
+    (( argument_count >= 2 )) || die "$option requires a value"
+}
+
+parse_args() {
+    while (($#)); do
+        case "$1" in
+            --prediction-prefix|--pred)
+                require_value "$1" "$#"
+                PREDICTION_PREFIX="$2"
+                shift 2
+                ;;
+            --genome-dir|--genome_dir)
+                require_value "$1" "$#"
+                GENOME_DIR="$2"
+                shift 2
+                ;;
+            --output-json)
+                require_value "$1" "$#"
+                OUTPUT_JSON="$2"
+                shift 2
+                ;;
+            --codon-table|--codon_table)
+                require_value "$1" "$#"
+                CODON_TABLE="$2"
+                shift 2
+                ;;
+            --confidence-threshold)
+                require_value "$1" "$#"
+                CONFIDENCE_THRESHOLD="$2"
+                shift 2
+                ;;
+            -h|--help)
+                usage 0
+                ;;
+            *)
+                printf 'Unknown argument: %s\n\n' "$1" >&2
+                usage 2
+                ;;
+        esac
+    done
+}
+
+main() {
+    parse_args "$@"
+    [[ -n "$PREDICTION_PREFIX" && -n "$GENOME_DIR" ]] || usage 2
+
+    [[ -s "$PREDICTION_PREFIX.gff" && -s "$PREDICTION_PREFIX.fna" ]] || \
+        die "missing prediction GFF/FNA: $PREDICTION_PREFIX.gff, $PREDICTION_PREFIX.fna"
+    [[ -s "$GENOME_DIR/genomic.gff" && -s "$GENOME_DIR/cds_from_genomic.fna" ]] || \
+        die "missing reference genomic.gff or cds_from_genomic.fna in $GENOME_DIR"
+
+    local output_json="${OUTPUT_JSON:-$(dirname -- "$PREDICTION_PREFIX")/evaluation_metrics.json}"
+    local evaluation_args=(
+        --pred "$PREDICTION_PREFIX"
+        --genome_dir "$GENOME_DIR"
+        --confidence-threshold "$CONFIDENCE_THRESHOLD"
+        --output-json "$output_json"
+    )
+    if [[ -n "$CODON_TABLE" ]]; then
+        [[ -f "$CODON_TABLE" ]] || die "codon table does not exist: $CODON_TABLE"
+        evaluation_args+=(--codon-table "$CODON_TABLE")
+    fi
+
+    echo "[evaluation] running $PREDICTION_PREFIX"
+    python "$ROOT_DIR/src/evaluation.py" "${evaluation_args[@]}"
+}
+
+main "$@"

annotation_pipeline/src/busco_evaluation.py

@@ -0,0 +1,212 @@
+"""Run BUSCO on the predicted proteins and retain only summary JSON."""
+
+import argparse
+from datetime import datetime, timezone
+import json
+from pathlib import Path
+import shutil
+import subprocess
+import tempfile
+
+from evaluation import iter_fasta_records
+
+
+BUSCO_RESULT_FIELDS = {
+    "busco_complete_percent": "Complete percentage",
+    "busco_complete_count": "Complete BUSCOs",
+    "busco_single_copy_percent": "Single copy percentage",
+    "busco_single_copy_count": "Single copy BUSCOs",
+    "busco_duplicated_percent": "Multi copy percentage",
+    "busco_duplicated_count": "Multi copy BUSCOs",
+    "busco_fragmented_percent": "Fragmented percentage",
+    "busco_fragmented_count": "Fragmented BUSCOs",
+    "busco_missing_percent": "Missing percentage",
+    "busco_missing_count": "Missing BUSCOs",
+}
+
+
+def parse_busco_summary(path, lineage):
+    """Extract stable key metrics from a BUSCO short_summary.json file."""
+    with Path(path).open(encoding="utf-8") as handle:
+        summary = json.load(handle)
+    results = summary["results"]
+    lineage_data = summary["lineage_dataset"]
+    versions = summary["versions"]
+    if lineage_data.get("name") != lineage:
+        raise ValueError(
+            f"BUSCO summary lineage is {lineage_data.get('name')!r}, "
+            f"expected {lineage!r}"
+        )
+    metrics = {
+        "busco_lineage": lineage,
+        "busco_marker_count": int(results["n_markers"]),
+        "busco_version": str(versions["busco"]),
+    }
+    for output_name, summary_name in BUSCO_RESULT_FIELDS.items():
+        metrics[output_name] = results[summary_name]
+    lineage_markers = int(lineage_data["number_of_buscos"])
+    if metrics["busco_marker_count"] != lineage_markers:
+        raise ValueError(
+            f"BUSCO marker count mismatch: results={metrics['busco_marker_count']}, "
+            f"lineage={lineage_markers}"
+        )
+    return metrics
+
+
+def run_busco_command(
+    executable,
+    input_faa,
+    lineage_path,
+    lineage,
+    output_directory,
+    workers,
+):
+    """Run BUSCO once and return parsed summary metrics."""
+    command = [
+        str(executable),
+        "-i",
+        str(input_faa),
+        "-l",
+        str(lineage_path),
+        "-o",
+        "busco",
+        "--out_path",
+        str(output_directory),
+        "-m",
+        "proteins",
+        "--offline",
+        "--opt-out-run-stats",
+        "-c",
+        str(workers),
+    ]
+    log_path = output_directory.parent / f"{output_directory.name}.log"
+    with log_path.open("w", encoding="utf-8") as log_handle:
+        completed = subprocess.run(
+            command,
+            stdout=log_handle,
+            stderr=subprocess.STDOUT,
+            check=False,
+        )
+    if completed.returncode != 0:
+        log_lines = log_path.read_text(encoding="utf-8", errors="replace").splitlines()
+        log_tail = "\n".join(log_lines[-40:])
+        raise RuntimeError(
+            f"BUSCO failed with exit code {completed.returncode}:\n{log_tail}"
+        )
+
+    summary_path = (
+        output_directory
+        / "busco"
+        / f"run_{lineage}"
+        / "short_summary.json"
+    )
+    if not summary_path.is_file():
+        raise RuntimeError(f"BUSCO did not create expected summary: {summary_path}")
+    return parse_busco_summary(summary_path, lineage)
+
+
+def write_json_atomic(path, payload):
+    output_path = Path(path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    temporary_path = output_path.with_name(output_path.name + ".tmp")
+    with temporary_path.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, sort_keys=True, allow_nan=False)
+        handle.write("\n")
+    temporary_path.replace(output_path)
+
+
+def count_proteins(path):
+    """Return the number of protein records, rejecting duplicate gene IDs."""
+    gene_ids = set()
+    for header, _ in iter_fasta_records(path):
+        gene_id = header.split(None, 1)[0]
+        if gene_id in gene_ids:
+            raise ValueError(f"duplicate protein FASTA gene ID {gene_id!r}")
+        gene_ids.add(gene_id)
+    if not gene_ids:
+        raise ValueError(f"{path} contains no proteins")
+    return len(gene_ids)
+
+
+def run_busco(
+    accession,
+    prediction_faa,
+    lineage,
+    lineage_path,
+    busco_executable,
+    workers,
+    output_json,
+    temporary_directory=None,
+):
+    """Run BUSCO on the predicted proteins and write one summary JSON."""
+    protein_count = count_proteins(prediction_faa)
+    with tempfile.TemporaryDirectory(
+        prefix=f"busco_{accession}_",
+        dir=temporary_directory,
+    ) as work_directory:
+        metrics = run_busco_command(
+            busco_executable,
+            prediction_faa,
+            lineage_path,
+            lineage,
+            Path(work_directory) / "run",
+            workers,
+        )
+    metrics["predicted_proteins"] = protein_count
+    print(
+        f"[busco] {accession} lineage={lineage} proteins={protein_count} "
+        f"complete={metrics['busco_complete_percent']:.1f}%",
+        flush=True,
+    )
+
+    payload = {
+        "accession": accession,
+        "generated_at": datetime.now(timezone.utc).isoformat(),
+        "lineage": lineage,
+        "metrics": metrics,
+    }
+    write_json_atomic(output_json, payload)
+    print(f"[busco] wrote {output_json}")
+    return payload
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--accession", required=True)
+    parser.add_argument("--faa", required=True, type=Path)
+    parser.add_argument("--lineage", required=True)
+    parser.add_argument("--lineage-path", required=True, type=Path)
+    parser.add_argument("--busco-executable", default="busco")
+    parser.add_argument("--workers", type=int, default=8)
+    parser.add_argument("--output-json", required=True, type=Path)
+    parser.add_argument("--temporary-directory", type=Path)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    if args.workers < 1:
+        raise SystemExit("--workers must be at least 1")
+    if shutil.which(args.busco_executable) is None:
+        raise SystemExit(f"BUSCO executable not found: {args.busco_executable}")
+    if not args.faa.is_file():
+        raise SystemExit(f"prediction FAA not found: {args.faa}")
+    if not args.lineage_path.is_dir():
+        raise SystemExit(f"lineage directory not found: {args.lineage_path}")
+    if args.temporary_directory is not None:
+        args.temporary_directory.mkdir(parents=True, exist_ok=True)
+
+    return run_busco(
+        accession=args.accession,
+        prediction_faa=args.faa,
+        lineage=args.lineage,
+        lineage_path=args.lineage_path,
+        busco_executable=args.busco_executable,
+        workers=args.workers,
+        output_json=args.output_json,
+        temporary_directory=args.temporary_directory,
+    )
+
+
+if __name__ == "__main__":
+    main()

annotation_pipeline/src/confidence_network.py

@@ -0,0 +1,232 @@
+"""Confidence network: one confidence score per decoded gene.
+
+The confidence network estimates, for every complete gene reconstructed by
+the Viterbi algorithm, the probability that its structure is exactly right.
+It reads the whole strand the gene lies on: the coding-probability track, and
+the role the Viterbi algorithm gave every base, exon, intron, or intergenic,
+neighbouring genes included. The two are combined into the log-odds that each
+base's role is right (the coding log-odds at exon bases, the non-coding
+log-odds elsewhere) plus the one-hot roles; a TCN with the same architecture
+as the transition network's (tcn.py) turns these four tracks into one feature
+vector per base. For each scored gene, the feature vectors at its exon bases
+are pooled by two attention heads and an MLP maps the pooled vectors to the
+confidence logit; the sigmoid of the logit is the confidence.
+
+``inference`` does this for a strand of any length without holding the whole
+strand on the device: a feature vector depends only on the strand within the
+network's context radius of its base, so only the windows of that radius
+around the scored exons are read, concatenated and processed in chunks the
+way the transition network reads a record (tcn.chunked_inference). Only
+complete genes are scored; a gene truncated by a record edge has no complete
+structure to judge, but its decoded bases still give roles to the strand for
+its neighbours.
+"""
+
+import numpy as np
+import torch
+from torch import nn
+
+from constants import NETWORK_CHUNK_BP
+from tcn import TCN, chunked_inference, coding_log_odds, load_weights
+
+CONFIDENCE_NETWORK_TYPE = "confidence_network"
+
+# Roles of the bases in the decoded structure, in channel order.
+ROLE_NAMES = ("exon", "intron", "intergenic")
+ROLE_EXON, ROLE_INTRON, ROLE_INTERGENIC = range(len(ROLE_NAMES))
+# Input tracks in channel order: the log-odds that the decoded role is right,
+# then one indicator track per role.
+INPUT_NAMES = ("role_log_odds",) + ROLE_NAMES
+
+
+def input_tracks(probabilities, roles):
+    """Return the ``[batch, inputs, bases]`` network input from coding probabilities and roles.
+
+    The first track is the log-odds that each base's decoded role is right:
+    the coding log-odds at exon bases and their negative, the non-coding
+    log-odds, at intron and intergenic bases. The remaining tracks are the
+    one-hot roles.
+    """
+    if probabilities.ndim != 2 or roles.shape != probabilities.shape:
+        raise ValueError("probabilities and roles must both have shape [batch, bases]")
+    log_odds = coding_log_odds(probabilities)
+    role_tracks = nn.functional.one_hot(roles, num_classes=len(ROLE_NAMES)).to(log_odds.dtype)
+    sign = role_tracks[..., ROLE_EXON] * 2 - 1
+    return torch.cat(((log_odds * sign)[:, None, :], role_tracks.transpose(1, 2)), dim=1)
+
+
+class ConfidenceNetwork(nn.Module):
+    """Score decoded genes from the strand's coding probabilities and decoded roles."""
+
+    def __init__(
+        self, channels=32, dilation_count=8,
+        attention_heads=2, attention_channels=32, readout_channels=32,
+    ):
+        super().__init__()
+        self.attention_heads = int(attention_heads)
+        self.encoder = TCN(len(INPUT_NAMES), channels, dilation_count)
+        self.context_radius = self.encoder.context_radius
+        # One attention logit per head and exon base, from that base's features.
+        self.attention = nn.Sequential(
+            nn.Linear(self.encoder.channels, attention_channels, bias=False),
+            nn.GELU(),
+            nn.Linear(attention_channels, self.attention_heads, bias=False),
+        )
+        # The pooled vectors of all heads, concatenated, give the confidence logit.
+        self.readout = nn.Sequential(
+            nn.Linear(self.attention_heads * self.encoder.channels, readout_channels),
+            nn.GELU(),
+            nn.Linear(readout_channels, 1),
+        )
+
+    def encode(self, probabilities, roles):
+        """Return the ``[bases, channels]`` features of a strand, or a chunk of one, given as ``[bases]`` tensors."""
+        return self.encoder(input_tracks(probabilities[None, :], roles[None, :]))[0]
+
+    def gene_logits(self, exon_features):
+        """Return one confidence logit per gene from the features at its exon bases.
+
+        ``exon_features`` holds one ``[exon_bases, channels]`` tensor per gene.
+        Each head gives every exon base an attention logit; the softmax over
+        the bases of one gene weights their feature vectors into one pooled
+        vector per head. Genes are pooled side by side with segment sums, so
+        no gene is padded to the length of another.
+        """
+        exon_features = tuple(exon_features)
+        if not exon_features:
+            raise ValueError("at least one gene is required")
+        for features in exon_features:
+            if features.ndim != 2 or not len(features):
+                raise ValueError("a gene must contain at least one exon base")
+
+        lengths = torch.as_tensor(
+            [len(features) for features in exon_features],
+            dtype=torch.long, device=exon_features[0].device,
+        )
+        features = torch.cat(exon_features)
+        # The gene every base belongs to, repeated for every head.
+        gene_of_base = torch.repeat_interleave(torch.arange(len(lengths), device=lengths.device), lengths)
+        gene_of_base_per_head = gene_of_base[:, None].expand(-1, self.attention_heads)
+
+        # Softmax within each gene: subtract the gene's maximum logit, exponentiate,
+        # divide by the gene's sum.
+        attention_logits = self.attention(features)
+        maximum = features.new_full((len(lengths), self.attention_heads), -torch.inf)
+        maximum.scatter_reduce_(0, gene_of_base_per_head, attention_logits, reduce="amax", include_self=True)
+        attention = torch.exp(attention_logits - maximum[gene_of_base])
+        normalization = features.new_zeros(len(lengths), self.attention_heads)
+        normalization.scatter_add_(0, gene_of_base_per_head, attention)
+        attention = attention / normalization[gene_of_base]
+
+        # Attention-weighted sum of the feature vectors of each gene, per head.
+        pooled = features.new_zeros(len(lengths), self.attention_heads, features.shape[1])
+        for head in range(self.attention_heads):
+            pooled[:, head].index_add_(0, gene_of_base, attention[:, head, None] * features)
+        return self.readout(pooled.flatten(1)).squeeze(1)
+
+    @torch.inference_mode()
+    def inference(self, probabilities, exons_by_gene, scored, chunk_bp=NETWORK_CHUNK_BP):
+        """Return the confidence of each scored gene of one strand.
+
+        ``probabilities`` is the strand's coding-probability track in decoder
+        orientation; ``exons_by_gene`` holds the inclusive exon intervals, in
+        transcript order, of every gene decoded on the strand; ``scored`` lists
+        the positions in ``exons_by_gene`` of the complete genes to score.
+        Every decoded gene gives roles to the strand. The confidences come back
+        as a float64 array in the order of ``scored``.
+        """
+        scored = list(scored)
+        if not scored:
+            return np.zeros(0, dtype=np.float64)
+        device = next(self.parameters()).device
+        probabilities = np.asarray(probabilities)
+        roles = strand_roles(exons_by_gene, len(probabilities))
+        scored_exons = [exons_by_gene[index] for index in scored]
+
+        # 1. Keep only the bases a scored exon can see: the window of one
+        #    context radius around every scored exon, merged, in strand order.
+        windows = exon_windows(scored_exons, len(probabilities), self.context_radius)
+        window_probabilities = np.concatenate([probabilities[start:stop] for start, stop in windows])
+        window_roles = np.concatenate([roles[start:stop] for start, stop in windows])
+        # 2. Where each gene's exon bases lie in the concatenated windows, in transcript order.
+        positions_by_gene = [
+            positions_in_windows(exon_positions(oriented_exons), windows) for oriented_exons in scored_exons
+        ]
+        # 3. Encode the concatenated windows in chunks; keep the features at the exon bases.
+        features = chunked_inference(
+            self.encode, self.context_radius, (window_probabilities, window_roles), device, chunk_bp,
+        )
+        exon_features = [torch.as_tensor(features[positions]).to(device) for positions in positions_by_gene]
+        # 4. Pool each gene's exon features into its confidence.
+        logits = self.gene_logits(exon_features).to(torch.float64)
+        return torch.sigmoid(logits).cpu().numpy()
+
+
+def load_confidence_network(path, device):
+    """Load the confidence network weights onto ``device`` in evaluation mode."""
+    checkpoint = load_weights(path)
+    if checkpoint.get("model_type") != CONFIDENCE_NETWORK_TYPE:
+        raise ValueError(f"{path}: not a confidence network checkpoint")
+    if tuple(checkpoint.get("input_names", ())) != INPUT_NAMES:
+        raise ValueError(f"{path}: unexpected input tracks")
+    if tuple(checkpoint.get("role_names", ())) != ROLE_NAMES:
+        raise ValueError(f"{path}: unexpected roles")
+    network = ConfidenceNetwork(
+        int(checkpoint["channels"]), int(checkpoint["dilation_count"]),
+        int(checkpoint["attention_heads"]), int(checkpoint["attention_channels"]),
+        int(checkpoint["readout_channels"]),
+    )
+    network.load_state_dict(checkpoint["state_dict"])
+    return network.to(device).eval()
+
+
+def strand_roles(exons_by_gene, record_length):
+    """Return the role of every base of a strand from the oriented exons of its decoded genes.
+
+    A base inside a gene's span is intron unless some gene has an exon there;
+    every other base is intergenic. Every decoded gene counts, complete or not.
+    """
+    roles = np.full(record_length, ROLE_INTERGENIC, dtype=np.int64)
+    for oriented_exons in exons_by_gene:
+        roles[oriented_exons[0][0]:oriented_exons[-1][1] + 1] = ROLE_INTRON
+    for oriented_exons in exons_by_gene:
+        for exon_start, exon_stop_inclusive in oriented_exons:
+            roles[exon_start:exon_stop_inclusive + 1] = ROLE_EXON
+    return roles
+
+
+def exon_positions(oriented_exons):
+    """Return the strand positions of a gene's exon bases, in transcript order."""
+    return np.concatenate([
+        np.arange(exon_start, exon_stop_inclusive + 1) for exon_start, exon_stop_inclusive in oriented_exons
+    ])
+
+
+def exon_windows(scored_exons, record_length, context_radius):
+    """Return the half-open windows of the strand within ``context_radius`` of a scored exon.
+
+    Every exon is widened by the radius on both sides and clipped to the
+    record; overlapping or touching windows are merged. Sorted along the
+    strand. Reading these windows and nothing else gives every exon base its
+    full context.
+    """
+    widened = sorted(
+        (max(0, exon_start - context_radius), min(record_length, exon_stop_inclusive + 1 + context_radius))
+        for oriented_exons in scored_exons
+        for exon_start, exon_stop_inclusive in oriented_exons
+    )
+    windows = []
+    for start, stop in widened:
+        if windows and start <= windows[-1][1]:
+            windows[-1][1] = max(windows[-1][1], stop)
+        else:
+            windows.append([start, stop])
+    return [(start, stop) for start, stop in windows]
+
+
+def positions_in_windows(positions, windows):
+    """Map sorted strand positions inside ``windows`` to positions in the concatenated windows."""
+    window_starts = np.asarray([start for start, _ in windows], dtype=np.int64)
+    concatenated_starts = np.concatenate(([0], np.cumsum([stop - start for start, stop in windows])[:-1]))
+    window = np.searchsorted(window_starts, positions, side="right") - 1
+    return positions - window_starts[window] + concatenated_starts[window]

annotation_pipeline/src/constants.py

@@ -0,0 +1,141 @@
+"""Shared constants for the Carbon-A annotation pipeline.
+
+Keep cross-module data contracts and tunable pipeline parameters here. Private
+implementation details, such as decoder state ids and evaluation regexes, stay
+with the modules that own them.
+"""
+
+import math
+from pathlib import Path
+
+
+PROJECT_DIRECTORY = Path(__file__).resolve().parent.parent
+
+
+# Model inference and output schema. The model reads windows of
+# CONTEXT_LENGTH tokens (six bases each); consecutive windows of a long record
+# overlap by OVERLAP_LENGTH tokens and overlapping probabilities are averaged.
+MODEL_NAME = "HuggingFaceBio/Carbon-A-1.2B"
+CONTEXT_LENGTH = 16384
+OVERLAP_LENGTH = 1024
+HEAD_NAMES = ("positive_strand", "negative_strand")
+RECORD_ID_COLUMN = "record_name"
+SEQUENCE_COLUMN = "sequence"
+POSITIVE_STRAND_PROBABILITY_COLUMN = "pred_prob_positive_strand_cds"
+NEGATIVE_STRAND_PROBABILITY_COLUMN = "pred_prob_negative_strand_cds"
+
+
+# Sequence alphabet and biological constraints.
+BASE_A = 0
+BASE_T = 1
+BASE_C = 2
+BASE_G = 3
+BASE_N = 4
+BASE_VALUES = {
+    "A": BASE_A,
+    "T": BASE_T,
+    "C": BASE_C,
+    "G": BASE_G,
+    "N": BASE_N,
+}
+COMPLEMENT_BASES = {
+    "A": "T",
+    "T": "A",
+    "C": "G",
+    "G": "C",
+    "N": "N",
+}
+START_CODONS = ("ATG",)
+STOP_CODONS = ("TAA", "TAG", "TGA")
+CODON_BASES = "ATCG"
+AMINO_ACID_SYMBOLS = frozenset("ACDEFGHIKLMNPQRSTVWY*")
+NON_NUCLEAR_ASSEMBLY_UNIT = "non-nuclear"
+
+
+# Viterbi and post-processing parameters.
+# Keep emission probabilities finite before taking logarithms.
+MIN_PROB = 1.0e-6
+MAX_PROB = 1.0 - MIN_PROB
+
+# Output gene models must contain at least 32 codons. A complete donor and
+# acceptor occupy four distinct intron bases.
+MIN_CDS_LENGTH_BP = 96
+MIN_INTERNAL_INTRON_LENGTH_BP = 4
+
+# Codon-table check. A codon table that does not match the genome splits genes
+# at codons the genome reads through or fuses genes at stop codons it lacks,
+# and the confidence network gives such genes almost no confidence. The check
+# warns when fewer than CODON_TABLE_WARNING_FRACTION of the complete genes
+# reach CONFIDENCE_THRESHOLD. With the shipped network, 17 genomes decoded
+# with their own table keep 67% (hagfish) to 98% of their complete genes at or
+# above the threshold; Tetrahymena decoded with the standard table keeps 1.4%
+# and yeast decoded with table 6 keeps 36%. A table that only drops a rarely
+# used stop codon escapes the check (Dictyostelium without TGA: 94%). Fewer
+# than CODON_TABLE_WARNING_MIN_GENES complete genes are too few to judge.
+CODON_TABLE_WARNING_FRACTION = 0.5
+CODON_TABLE_WARNING_MIN_GENES = 100
+# A failed check is recorded where a batch run will see it: appended to this
+# file in the parent of the output directory, and in the output directory's
+# name (run_*.sh share the suffix).
+WARNINGS_FILE_NAME = "WARNINGS.txt"
+WARNING_DIRECTORY_SUFFIX = "_WARNING"
+
+# Intron classes are scored relative to the 99% GT-AG baseline.
+GT_AG_INTRON_LOG_PRIOR = math.log(0.990 / 0.99)
+GC_AG_INTRON_LOG_PRIOR = math.log(0.009 / 0.99)
+AT_AC_INTRON_LOG_PRIOR = math.log(0.001 / 0.99)
+
+# The two networks of the annotation pipeline. The transition network reads
+# the coding-probability track and scores the four permitted transitions at
+# every base; the confidence network reads the track together with the
+# decoded roles of the whole strand and scores every complete gene. Both
+# read coding probabilities clipped to NETWORK_INPUT_MIN_PROB as log-odds.
+WEIGHTS_DIRECTORY = PROJECT_DIRECTORY / "weights"
+TRANSITION_NETWORK_PATH = WEIGHTS_DIRECTORY / "transition_network.safetensors"
+CONFIDENCE_NETWORK_PATH = WEIGHTS_DIRECTORY / "confidence_network.safetensors"
+# Per-base transition probability tracks handed to the Viterbi backends, in
+# this column order. They are stored as log-odds against staying in the
+# backends' fixed-point units; the native backends and the Numba reference
+# share the layout.
+TRANSITION_NAMES = ("enter_gene", "leave_gene", "enter_intron", "leave_intron")
+TRANSITION_COUNT = len(TRANSITION_NAMES)
+ENTER_GENE, LEAVE_GENE, ENTER_INTRON, LEAVE_INTRON = range(TRANSITION_COUNT)
+NETWORK_INPUT_MIN_PROB = 1.0e-4
+# Both networks read long tracks in chunks of this many bases, each evaluated
+# with the network's full context on both sides (tcn.chunked_inference), so a
+# chromosome never has to fit in GPU memory at once. (The post-processing
+# driver separately groups whole records into host-side chunks, chunk_mbp.)
+NETWORK_CHUNK_BP = 262_144
+# Native Viterbi implementations accumulate int32 fixed-point scores.
+# Transition probabilities are exchanged as log-odds in the same units: nats
+# multiplied by this scale and rounded.
+FIXED_POINT_SCORE_SCALE = 16384.0
+
+# Boundary introns model genes truncated at a record edge. A finite limit
+# prevents an incomplete path from hiding a long region as one intron.
+MAX_BOUNDARY_INTRON_LENGTH_BP = 20_000
+
+# Records shorter than this are skipped. Host-side chunks can exceed their
+# target size because a sequence record is never split.
+MIN_RECORD_LENGTH_BP = 100_000
+CHUNK_SIZE_MBP = 25
+PARQUET_BATCH_ROWS = 1
+CODON_TABLE_PATH = PROJECT_DIRECTORY / "codon_tables" / "standard.json"
+
+
+# GPU Viterbi backend memory and batching parameters.
+GPU_SEGMENT_LENGTH = 2000
+GPU_MEMORY_FRACTION = 0.70
+GPU_MIN_BATCH_BASES = 20_000_000
+# The CUDA preprocessing kernel indexes packed bases with int32. Record offsets
+# remain int64, so this limits a sub-batch rather than the input file.
+GPU_MAX_BATCH_BASES = (1 << 31) - 1
+# Covers packed inputs, per-base scores and paths, the four int32 transition
+# score tracks, segment matrices, and allocation headroom while CUDA windows
+# the backpointer array internally.
+GPU_ESTIMATED_BYTES_PER_BASE = 34
+
+
+# Gene-level results count the predicted genes whose confidence reaches this
+# threshold; genes without a confidence are always counted.
+CONFIDENCE_THRESHOLD = 0.1

annotation_pipeline/src/evaluation.py

@@ -0,0 +1,1093 @@
+"""Evaluate predicted CDS annotations against an NCBI reference assembly.
+
+When ``sequence_report.jsonl`` is available, evaluation prefers assembled
+chromosomes, then assembled linkage groups, then chromosome scaffolds; without
+the report every record is evaluated. Six complementary metrics are reported:
+
+* gene: one-to-one exact CDS-sequence matches on the same genomic record;
+* protein: gene exact matches plus additional one-to-one exact matches after
+  translating unmatched predicted and reference CDS with the same codon table;
+* exon: exact merged CDS intervals from the longest CDS isoform per gene;
+* 5' boundary and 3' boundary: individual CDS boundaries in transcript direction;
+* nucleotide: base-pair overlap against those reference intervals.
+
+Gene and protein metrics count the predicted genes whose ``confidence``
+reaches the confidence threshold (default 0.1); genes without a confidence are
+always counted, and a threshold of 0 counts every predicted gene. Exon,
+boundary, and nucleotide metrics describe the whole predicted annotation.
+Evaluation also reports AUROC, non-interpolated AUPRC, and point-biserial
+correlation of the confidence against the one-to-one gene exact-match label,
+together with exact and non-exact group means.
+
+The 5' and 3' boundary metrics follow transcript direction, so their meanings are
+reversed on the negative genomic strand. Interval arithmetic keeps memory
+proportional to the number of CDS intervals rather than genome length.
+
+Example:
+
+    python src/evaluation.py --pred results/GCF_000146045.2/prediction \
+        --genome_dir ncbi_dataset/data/GCF_000146045.2
+"""
+
+import argparse
+from collections import defaultdict
+from datetime import datetime, timezone
+import json
+import math
+import os
+from pathlib import Path
+import re
+import sys
+
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from constants import CONFIDENCE_THRESHOLD
+from utils import load_codon_table, load_nuclear_evaluation_accessions
+
+
+STRANDS = ("+", "-")
+REFERENCE_CDS_FEATURES = frozenset(("mRNA", "transcript", "CDS"))
+METRIC_NAMES = (
+    "gene",
+    "protein",
+    "exon",
+    "five_prime_boundary",
+    "three_prime_boundary",
+    "nucleotide",
+)
+INTERVAL_METRIC_NAMES = (
+    "exon",
+    "five_prime_boundary",
+    "three_prime_boundary",
+    "nucleotide",
+)
+REFERENCE_RECORD_PATTERN = re.compile(r"^lcl\|(.+?)_cds_")
+PREDICTED_RECORD_PATTERN = re.compile(r"(?:^|\s)record_id=(\S+)")
+PREDICTED_FRAME_PATTERN = re.compile(r"(?:^|\s)phase=([012])(?:\s|$)")
+REFERENCE_FRAME_PATTERN = re.compile(r"\[frame=([123])\]")
+GENE_ID_PATTERN = re.compile(
+    r"\[db_xref=(?:[^,\]]+,)*GeneID:(\d+)(?:,|\])"
+)
+LOCUS_TAG_PATTERN = re.compile(r"\[locus_tag=([^\]]+)\]")
+GENE_NAME_PATTERN = re.compile(r"\[gene=([^\]]+)\]")
+CONFIDENCE_NAME = "confidence"
+
+
+def iter_gff_rows(path):
+    """Yield the nine tab-separated fields from each GFF3 feature row."""
+    with open(path, encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, 1):
+            if line.startswith("#") or line.isspace():
+                continue
+            fields = line.rstrip("\n").split("\t")
+            if len(fields) != 9:
+                raise ValueError(
+                    f"{path}:{line_number}: expected 9 GFF3 fields, "
+                    f"found {len(fields)}"
+                )
+            yield fields
+
+
+def gff_attribute_value(attribute_text, name):
+    """Return one GFF3 attribute without parsing unrelated attributes."""
+    prefix = name + "="
+    if attribute_text.startswith(prefix):
+        value_start = len(prefix)
+    else:
+        marker = ";" + prefix
+        marker_start = attribute_text.find(marker)
+        if marker_start == -1:
+            return None
+        value_start = marker_start + len(marker)
+
+    value_end = attribute_text.find(";", value_start)
+    if value_end == -1:
+        return attribute_text[value_start:]
+    return attribute_text[value_start:value_end]
+
+
+def iter_selected_gff_features(path, feature_types, retained_records=None):
+    """Yield selected GFF3 features, optionally restricted by record id.
+
+    NCBI GFF files contain many gene, exon, and regulatory rows that CDS
+    evaluation never uses. Reading the first three columns lets those rows be
+    rejected before splitting their long attribute columns.
+    """
+    with open(path, encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, 1):
+            if line.startswith("#") or line.isspace():
+                continue
+
+            leading_fields = line.split("\t", 3)
+            if len(leading_fields) != 4:
+                raise ValueError(
+                    f"{path}:{line_number}: expected 9 GFF3 fields, "
+                    f"found {len(leading_fields)}"
+                )
+            record_id, _, feature_type, remaining_text = leading_fields
+            if feature_type not in feature_types:
+                continue
+            if retained_records is not None and record_id not in retained_records:
+                continue
+
+            remaining_fields = remaining_text.rstrip("\n").split("\t", 5)
+            if len(remaining_fields) != 6:
+                raise ValueError(
+                    f"{path}:{line_number}: expected 9 GFF3 fields, "
+                    f"found {3 + len(remaining_fields)}"
+                )
+            start, end, _, strand, _, attribute_text = remaining_fields
+            yield record_id, feature_type, start, end, strand, attribute_text
+
+
+def parse_predicted_cds_tracks(path):
+    """Return predicted CDS intervals keyed by ``(record_id, strand)``."""
+    tracks = defaultdict(list)
+    for fields in iter_gff_rows(path):
+        record_id, _, feature_type, start, end, _, strand, _, _ = fields
+        if feature_type != "CDS":
+            continue
+        if strand not in STRANDS:
+            raise ValueError(f"unsupported CDS strand {strand!r} in {path}")
+        tracks[(record_id, strand)].append((int(start) - 1, int(end) - 1))
+    return dict(tracks)
+
+
+def parse_reference_cds_tracks(path, retained_records=None):
+    """Return CDS intervals from the longest CDS isoform of each reference gene."""
+    transcript_gene = {}
+    transcript_cds = defaultdict(list)
+
+    for feature in iter_selected_gff_features(
+        path,
+        REFERENCE_CDS_FEATURES,
+        retained_records,
+    ):
+        record_id, feature_type, start, end, strand, attribute_text = feature
+        if feature_type in ("mRNA", "transcript"):
+            transcript_id = gff_attribute_value(attribute_text, "ID")
+            parent = gff_attribute_value(attribute_text, "Parent")
+            if transcript_id and parent:
+                transcript_gene[(record_id, transcript_id)] = parent.split(",", 1)[0]
+            continue
+
+        if strand not in STRANDS:
+            raise ValueError(f"unsupported CDS strand {strand!r} in {path}")
+        parent = gff_attribute_value(attribute_text, "Parent")
+        for transcript_id in (parent or "").split(","):
+            if transcript_id:
+                transcript_cds[(record_id, transcript_id)].append(
+                    (int(start) - 1, int(end) - 1, strand)
+                )
+
+    transcripts_by_gene = defaultdict(dict)
+    for (record_id, transcript_id), intervals in transcript_cds.items():
+        gene_id = transcript_gene.get((record_id, transcript_id), transcript_id)
+        transcripts_by_gene[(record_id, gene_id)][transcript_id] = intervals
+
+    tracks = defaultdict(list)
+    for (record_id, _), transcripts in transcripts_by_gene.items():
+        longest_transcript = max(
+            transcripts,
+            key=lambda transcript_id: sum(
+                end - start + 1
+                for start, end, _ in transcripts[transcript_id]
+            ),
+        )
+        intervals = transcripts[longest_transcript]
+        # NCBI trans-spliced transcripts can contain CDS pieces on both strands.
+        # Each piece remains on its annotated genomic strand for interval scoring.
+        for start, end, strand in intervals:
+            tracks[(record_id, strand)].append((start, end))
+    return dict(tracks)
+
+
+def iter_fasta_records(path):
+    """Yield ``(header_without_marker, sequence)`` from a FASTA file."""
+    header = None
+    sequence_lines = []
+    with open(path, encoding="utf-8") as handle:
+        for line in handle:
+            if line.startswith(">"):
+                if header is not None:
+                    yield header, "".join(sequence_lines)
+                header = line[1:].strip()
+                sequence_lines = []
+            else:
+                sequence_lines.append(line.strip())
+    if header is not None:
+        yield header, "".join(sequence_lines)
+
+
+def parse_predicted_cds_fasta(path):
+    """Return every predicted CDS as ``(record_id, sequence)``.
+
+    A list is intentional: two identical predictions remain two predictions and
+    therefore participate independently in precision accounting.
+    """
+    return [
+        (record_id, sequence)
+        for _, record_id, sequence in parse_predicted_cds_records(path)
+    ]
+
+
+def parse_predicted_cds_records(path):
+    """Return predicted CDS records as ``(gene_id, record_id, sequence)``."""
+    return [
+        (gene_id, record_id, sequence)
+        for gene_id, record_id, sequence, _ in (
+            parse_predicted_cds_records_with_frames(path)
+        )
+    ]
+
+
+def parse_predicted_cds_records_with_frames(path):
+    """Return predicted CDS records with their zero-based translation frame."""
+    predictions = []
+    for header, sequence in iter_fasta_records(path):
+        record_match = PREDICTED_RECORD_PATTERN.search(header)
+        if record_match is None:
+            raise ValueError(f"predicted FASTA header lacks record_id: {header}")
+        frame_match = PREDICTED_FRAME_PATTERN.search(header)
+        frame = int(frame_match.group(1)) if frame_match else 0
+        gene_id = header.split(None, 1)[0]
+        predictions.append((
+            gene_id,
+            record_match.group(1),
+            sequence.upper(),
+            frame,
+        ))
+    return predictions
+
+
+def translate_cds_for_evaluation(sequence, frame, codon_to_amino_acid):
+    """Translate one CDS for exact-match evaluation.
+
+    Incomplete trailing codons are ignored, ambiguous codons become ``X``, and
+    one terminal stop is removed. Internal stops remain in the sequence so an
+    unusual reference CDS can still be compared instead of aborting evaluation.
+    """
+    sequence = sequence.upper()
+    codon_end = len(sequence) - ((len(sequence) - frame) % 3)
+    protein = "".join(
+        codon_to_amino_acid.get(sequence[start:start + 3], "X")
+        for start in range(frame, codon_end, 3)
+    )
+    if protein.endswith("*"):
+        protein = protein[:-1]
+    return protein
+
+
+def parse_predicted_translated_proteins(path, codon_to_amino_acid):
+    """Translate predicted CDS records as ``(gene_id, record_id, protein)``."""
+    return [
+        (
+            gene_id,
+            record_id,
+            translate_cds_for_evaluation(sequence, frame, codon_to_amino_acid),
+        )
+        for gene_id, record_id, sequence, frame in (
+            parse_predicted_cds_records_with_frames(path)
+        )
+    ]
+
+
+def parse_predicted_gene_scores(path, score_name=CONFIDENCE_NAME):
+    """Return one optional gene-level score keyed by predicted GFF3 gene ID."""
+    scores = {}
+    for fields in iter_gff_rows(path):
+        _, _, feature_type, _, _, _, _, _, attribute_text = fields
+        if feature_type != "gene":
+            continue
+        gene_id = gff_attribute_value(attribute_text, "ID")
+        if gene_id is None:
+            raise ValueError(
+                f"predicted GFF3 gene lacks an ID in {path}"
+            )
+        if gene_id in scores:
+            raise ValueError(f"duplicate predicted gene ID {gene_id!r} in {path}")
+        score_text = gff_attribute_value(attribute_text, score_name)
+        if score_text is None:
+            scores[gene_id] = None
+            continue
+        if score_text in ("NA", "."):
+            scores[gene_id] = None
+            continue
+        try:
+            score = float(score_text)
+        except ValueError as exc:
+            raise ValueError(
+                f"invalid {score_name}={score_text!r} for {gene_id!r} in {path}"
+            ) from exc
+        if not math.isfinite(score) or not 0.0 <= score <= 1.0:
+            raise ValueError(
+                f"{score_name} for {gene_id!r} must be between 0 and 1, "
+                f"found {score_text!r}"
+            )
+        scores[gene_id] = score
+    return scores
+
+
+def reference_gene_id(header):
+    """Return the most stable gene identifier available in an NCBI CDS header."""
+    match = GENE_ID_PATTERN.search(header)
+    if match:
+        return "GeneID:" + match.group(1)
+    match = LOCUS_TAG_PATTERN.search(header)
+    if match:
+        return "locus_tag:" + match.group(1)
+    match = GENE_NAME_PATTERN.search(header)
+    if match:
+        return "gene:" + match.group(1)
+    return "cds:" + header.split()[0]
+
+
+def parse_reference_cds_fasta(path):
+    """Return reference isoform sequences keyed by ``(record_id, gene_id)``."""
+    return {
+        gene_key: [sequence for sequence, _ in isoforms]
+        for gene_key, isoforms in parse_reference_cds_fasta_with_frames(path).items()
+    }
+
+
+def parse_reference_cds_fasta_with_frames(path):
+    """Return reference CDS and zero-based frame by genomic record and gene."""
+    genes = defaultdict(list)
+    for header, sequence in iter_fasta_records(path):
+        record_match = REFERENCE_RECORD_PATTERN.search(header)
+        if record_match is None:
+            raise ValueError(f"unrecognized NCBI CDS FASTA header: {header}")
+        frame_match = REFERENCE_FRAME_PATTERN.search(header)
+        frame = int(frame_match.group(1)) - 1 if frame_match else 0
+        key = (record_match.group(1), reference_gene_id(header))
+        genes[key].append((sequence.upper(), frame))
+    return dict(genes)
+
+
+def parse_reference_translated_proteins(path, codon_to_amino_acid):
+    """Translate reference CDS isoforms by ``(record_id, gene_id)``."""
+    return {
+        gene_key: [
+            translate_cds_for_evaluation(sequence, frame, codon_to_amino_acid)
+            for sequence, frame in isoforms
+        ]
+        for gene_key, isoforms in parse_reference_cds_fasta_with_frames(path).items()
+    }
+
+
+def merge_intervals(intervals):
+    """Return the union of inclusive intervals in coordinate order."""
+    if not intervals:
+        return []
+    ordered = sorted(intervals)
+    if any(start > end for start, end in ordered):
+        raise ValueError(f"invalid interval list: {ordered}")
+    merged = [list(ordered[0])]
+    for start, end in ordered[1:]:
+        previous = merged[-1]
+        if start <= previous[1] + 1:
+            previous[1] = max(previous[1], end)
+        else:
+            merged.append([start, end])
+    return [(start, end) for start, end in merged]
+
+
+def merge_tracks(tracks):
+    """Merge overlapping or adjacent CDS intervals independently per track."""
+    return {track: merge_intervals(intervals) for track, intervals in tracks.items()}
+
+
+def interval_overlap_length(left, right):
+    """Return the number of bases covered by both merged interval lists."""
+    overlap = 0
+    left_index = 0
+    right_index = 0
+    while left_index < len(left) and right_index < len(right):
+        left_start, left_end = left[left_index]
+        right_start, right_end = right[right_index]
+        overlap += max(
+            0,
+            min(left_end, right_end) - max(left_start, right_start) + 1,
+        )
+        if left_end < right_end:
+            left_index += 1
+        else:
+            right_index += 1
+    return overlap
+
+
+def add_set_metric_counts(counts, metric_name, predicted_items, reference_items):
+    """Accumulate TP/FP/FN counts for two sets."""
+    true_positives = len(predicted_items & reference_items)
+    counts[f"{metric_name}_tp"] += true_positives
+    counts[f"{metric_name}_fp"] += len(predicted_items) - true_positives
+    counts[f"{metric_name}_fn"] += len(reference_items) - true_positives
+
+
+def interval_metric_counts(predicted_tracks, reference_tracks):
+    """Count exon, boundary, and nucleotide matches."""
+    predicted_tracks = merge_tracks(predicted_tracks)
+    reference_tracks = merge_tracks(reference_tracks)
+    counts = defaultdict(int)
+    for metric_name in INTERVAL_METRIC_NAMES:
+        for suffix in ("tp", "fp", "fn"):
+            counts[f"{metric_name}_{suffix}"] = 0
+
+    for track in set(predicted_tracks) | set(reference_tracks):
+        _, strand = track
+        predicted_exons = predicted_tracks.get(track, ())
+        reference_exons = reference_tracks.get(track, ())
+
+        predicted_bp = sum(end - start + 1 for start, end in predicted_exons)
+        reference_bp = sum(end - start + 1 for start, end in reference_exons)
+        overlap_bp = interval_overlap_length(predicted_exons, reference_exons)
+        counts["nucleotide_tp"] += overlap_bp
+        counts["nucleotide_fp"] += predicted_bp - overlap_bp
+        counts["nucleotide_fn"] += reference_bp - overlap_bp
+
+        predicted_left_boundaries = {start for start, _ in predicted_exons}
+        predicted_right_boundaries = {end for _, end in predicted_exons}
+        reference_left_boundaries = {start for start, _ in reference_exons}
+        reference_right_boundaries = {end for _, end in reference_exons}
+
+        if strand == "+":
+            predicted_five_prime_boundaries = predicted_left_boundaries
+            predicted_three_prime_boundaries = predicted_right_boundaries
+            reference_five_prime_boundaries = reference_left_boundaries
+            reference_three_prime_boundaries = reference_right_boundaries
+        else:
+            predicted_five_prime_boundaries = predicted_right_boundaries
+            predicted_three_prime_boundaries = predicted_left_boundaries
+            reference_five_prime_boundaries = reference_right_boundaries
+            reference_three_prime_boundaries = reference_left_boundaries
+
+        add_set_metric_counts(
+            counts,
+            "five_prime_boundary",
+            predicted_five_prime_boundaries,
+            reference_five_prime_boundaries,
+        )
+        add_set_metric_counts(
+            counts,
+            "three_prime_boundary",
+            predicted_three_prime_boundaries,
+            reference_three_prime_boundaries,
+        )
+        add_set_metric_counts(
+            counts,
+            "exon",
+            set(predicted_exons),
+            set(reference_exons),
+        )
+
+    return dict(counts)
+
+
+def gene_exact_match_assignments(predicted_cds, reference_genes):
+    """Return one-to-one exact matches keyed by prediction index.
+
+    A reference gene may match any one of its annotated isoforms. The augmenting
+    path prevents one prediction from recalling several genes with an identical
+    CDS and prevents duplicate predictions from matching the same gene twice.
+    """
+    candidate_genes = defaultdict(list)
+    for gene_key, isoforms in reference_genes.items():
+        record_id, _ = gene_key
+        for sequence in set(isoforms):
+            candidate_genes[(record_id, sequence)].append(gene_key)
+
+    matched_prediction_by_gene = {}
+    matched_gene_by_prediction = {}
+
+    def match_prediction(start_prediction):
+        """Find and apply one augmenting path without recursive calls."""
+        queued_predictions = [start_prediction]
+        visited_predictions = {start_prediction}
+        visited_genes = set()
+        parent_prediction_by_gene = {}
+        queue_index = 0
+
+        while queue_index < len(queued_predictions):
+            prediction_index = queued_predictions[queue_index]
+            queue_index += 1
+            prediction = predicted_cds[prediction_index]
+
+            for gene_key in candidate_genes.get(prediction, ()):
+                if gene_key in visited_genes:
+                    continue
+                visited_genes.add(gene_key)
+                parent_prediction_by_gene[gene_key] = prediction_index
+
+                previous_prediction = matched_prediction_by_gene.get(gene_key)
+                if previous_prediction is None:
+                    current_gene = gene_key
+                    while current_gene is not None:
+                        current_prediction = parent_prediction_by_gene[current_gene]
+                        previous_gene = matched_gene_by_prediction.get(
+                            current_prediction
+                        )
+                        matched_prediction_by_gene[current_gene] = current_prediction
+                        matched_gene_by_prediction[current_prediction] = current_gene
+                        current_gene = previous_gene
+                    return True
+
+                if previous_prediction not in visited_predictions:
+                    visited_predictions.add(previous_prediction)
+                    queued_predictions.append(previous_prediction)
+        return False
+
+    for prediction_index in range(len(predicted_cds)):
+        match_prediction(prediction_index)
+    return matched_gene_by_prediction
+
+
+def gene_metric_counts(
+    predicted_cds,
+    reference_genes,
+    matched_gene_by_prediction=None,
+):
+    """Count one-to-one exact CDS matches between predictions and genes."""
+    if matched_gene_by_prediction is None:
+        matched_gene_by_prediction = gene_exact_match_assignments(
+            predicted_cds,
+            reference_genes,
+        )
+    matched_genes = len(matched_gene_by_prediction)
+    return {
+        "gene_tp": matched_genes,
+        "gene_fp": len(predicted_cds) - matched_genes,
+        "gene_fn": len(reference_genes) - matched_genes,
+        "predicted_genes": len(predicted_cds),
+        "reference_genes": len(reference_genes),
+    }
+
+
+def protein_metric_counts(
+    predicted_proteins,
+    reference_proteins,
+    inherited_gene_matches=None,
+):
+    """Count protein matches, retaining every exact CDS gene match.
+
+    Translation frames can be uncertain for partial CDS. Exact nucleotide CDS
+    matches supplied in ``inherited_gene_matches`` therefore remain matches at
+    protein level even if their independently annotated frames differ. Protein
+    matching is then applied only to the still-unmatched predictions and genes,
+    allowing synonymous CDS differences to add matches without double counting.
+    """
+    if inherited_gene_matches is None:
+        matched_gene_by_prediction = gene_exact_match_assignments(
+            predicted_proteins,
+            reference_proteins,
+        )
+    else:
+        matched_gene_by_prediction = dict(inherited_gene_matches)
+        matched_reference_genes = set(matched_gene_by_prediction.values())
+        if len(matched_reference_genes) != len(matched_gene_by_prediction):
+            raise ValueError("inherited gene matches must be one-to-one")
+        if any(
+            prediction_index < 0 or prediction_index >= len(predicted_proteins)
+            for prediction_index in matched_gene_by_prediction
+        ):
+            raise ValueError("inherited gene match has an invalid prediction index")
+        if not matched_reference_genes.issubset(reference_proteins):
+            raise ValueError("inherited gene match is absent from protein references")
+
+        unmatched_prediction_indices = [
+            prediction_index
+            for prediction_index in range(len(predicted_proteins))
+            if prediction_index not in matched_gene_by_prediction
+        ]
+        unmatched_reference_proteins = {
+            gene_key: isoforms
+            for gene_key, isoforms in reference_proteins.items()
+            if gene_key not in matched_reference_genes
+        }
+        additional_matches = gene_exact_match_assignments(
+            [
+                predicted_proteins[prediction_index]
+                for prediction_index in unmatched_prediction_indices
+            ],
+            unmatched_reference_proteins,
+        )
+        matched_gene_by_prediction.update({
+            unmatched_prediction_indices[local_prediction_index]: gene_key
+            for local_prediction_index, gene_key in additional_matches.items()
+        })
+
+    matched_genes = len(matched_gene_by_prediction)
+    return {
+        "protein_tp": matched_genes,
+        "protein_fp": len(predicted_proteins) - matched_genes,
+        "protein_fn": len(reference_proteins) - matched_genes,
+        "predicted_proteins": len(predicted_proteins),
+        "reference_proteins": len(reference_proteins),
+    }
+
+
+def pearson_correlation(left_values, right_values):
+    """Return Pearson correlation, or ``None`` when either side is constant."""
+    if len(left_values) != len(right_values):
+        raise ValueError("correlation vectors must have the same length")
+    if len(left_values) < 2:
+        return None
+
+    left_mean = math.fsum(left_values) / len(left_values)
+    right_mean = math.fsum(right_values) / len(right_values)
+    left_deviations = [value - left_mean for value in left_values]
+    right_deviations = [value - right_mean for value in right_values]
+    left_sum_squares = math.fsum(value * value for value in left_deviations)
+    right_sum_squares = math.fsum(value * value for value in right_deviations)
+    if left_sum_squares == 0.0 or right_sum_squares == 0.0:
+        return None
+
+    cross_product = math.fsum(
+        left * right
+        for left, right in zip(left_deviations, right_deviations)
+    )
+    correlation = cross_product / math.sqrt(
+        left_sum_squares * right_sum_squares
+    )
+    return max(-1.0, min(1.0, correlation))
+
+
+def binary_ranking_auroc_auprc(scores, labels):
+    """Return tie-aware AUROC and stepwise AUPRC for binary labels.
+
+    Score ties are processed as a group. This gives tied positive/negative pairs
+    half credit in AUROC and prevents input order from affecting either metric.
+    AUPRC uses the average-precision (non-interpolated step) convention. Both
+    metrics are undefined when either class is absent.
+    """
+    if len(scores) != len(labels):
+        raise ValueError("ranking score and label vectors must have the same length")
+    if any(label not in (0, 1, False, True) for label in labels):
+        raise ValueError("ranking labels must be binary")
+
+    positive_count = sum(bool(label) for label in labels)
+    negative_count = len(labels) - positive_count
+    if positive_count == 0 or negative_count == 0:
+        return None, None
+
+    ranked = sorted(
+        zip(scores, (bool(label) for label in labels)),
+        key=lambda item: item[0],
+        reverse=True,
+    )
+    true_positives = 0
+    false_positives = 0
+    previous_recall = 0.0
+    auroc_area = 0.0
+    auprc_area = 0.0
+    group_start = 0
+    while group_start < len(ranked):
+        score = ranked[group_start][0]
+        group_end = group_start
+        group_positives = 0
+        while group_end < len(ranked) and ranked[group_end][0] == score:
+            group_positives += int(ranked[group_end][1])
+            group_end += 1
+        group_negatives = group_end - group_start - group_positives
+
+        next_true_positives = true_positives + group_positives
+        next_false_positives = false_positives + group_negatives
+        false_positive_rate_increment = group_negatives / negative_count
+        auroc_area += false_positive_rate_increment * (
+            true_positives + next_true_positives
+        ) / (2.0 * positive_count)
+
+        recall = next_true_positives / positive_count
+        precision = next_true_positives / (
+            next_true_positives + next_false_positives
+        )
+        auprc_area += (recall - previous_recall) * precision
+
+        true_positives = next_true_positives
+        false_positives = next_false_positives
+        previous_recall = recall
+        group_start = group_end
+
+    return auroc_area, auprc_area
+
+
+def confidence_exact_match_metrics(
+    confidence_values,
+    matched_gene_by_prediction,
+):
+    """Summarize association between confidence and exact gene matches."""
+    scored_values = []
+    exact_labels = []
+    exact_scores = []
+    nonexact_scores = []
+    for prediction_index, score in enumerate(confidence_values):
+        if score is None:
+            continue
+        is_exact = prediction_index in matched_gene_by_prediction
+        scored_values.append(score)
+        exact_labels.append(float(is_exact))
+        if is_exact:
+            exact_scores.append(score)
+        else:
+            nonexact_scores.append(score)
+
+    auroc, auprc = binary_ranking_auroc_auprc(scored_values, exact_labels)
+    return {
+        "confidence_scored_genes": len(scored_values),
+        "confidence_missing_genes": (
+            len(confidence_values) - len(scored_values)
+        ),
+        "confidence_mean": (
+            math.fsum(scored_values) / len(scored_values)
+            if scored_values
+            else None
+        ),
+        "confidence_exact_genes": len(exact_scores),
+        "confidence_nonexact_genes": len(nonexact_scores),
+        "confidence_exact_mean": (
+            math.fsum(exact_scores) / len(exact_scores)
+            if exact_scores
+            else None
+        ),
+        "confidence_nonexact_mean": (
+            math.fsum(nonexact_scores) / len(nonexact_scores)
+            if nonexact_scores
+            else None
+        ),
+        "confidence_exact_match_correlation": pearson_correlation(
+            scored_values,
+            exact_labels,
+        ),
+        "confidence_exact_match_auroc": auroc,
+        "confidence_exact_match_auprc": auprc,
+    }
+
+
+def precision_recall_f1(true_positives, false_positives, false_negatives):
+    """Return precision, recall, and F1 for integer match counts."""
+    precision_denominator = true_positives + false_positives
+    recall_denominator = true_positives + false_negatives
+    precision = (
+        true_positives / precision_denominator
+        if precision_denominator
+        else 0.0
+    )
+    recall = (
+        true_positives / recall_denominator
+        if recall_denominator
+        else 0.0
+    )
+    f1 = (
+        2 * precision * recall / (precision + recall)
+        if precision + recall
+        else 0.0
+    )
+    return precision, recall, f1
+
+
+def add_metric_rates(metrics):
+    """Add precision, recall, and F1 values for every metric count group."""
+    for metric_name in METRIC_NAMES:
+        precision, recall, f1 = precision_recall_f1(
+            metrics.get(f"{metric_name}_tp", 0),
+            metrics.get(f"{metric_name}_fp", 0),
+            metrics.get(f"{metric_name}_fn", 0),
+        )
+        metrics[f"{metric_name}_precision"] = precision
+        metrics[f"{metric_name}_recall"] = recall
+        metrics[f"{metric_name}_f1"] = f1
+
+
+def filter_tracks(tracks, retained_records):
+    """Restrict interval tracks to a record set, or return all tracks for None."""
+    if retained_records is None:
+        return tracks
+    return {
+        track: intervals
+        for track, intervals in tracks.items()
+        if track[0] in retained_records
+    }
+
+
+def filter_gene_inputs(predicted_cds, reference_genes, retained_records):
+    """Restrict gene-level inputs to a record set."""
+    if retained_records is None:
+        return predicted_cds, reference_genes
+    return (
+        [
+            prediction
+            for prediction in predicted_cds
+            if prediction[0] in retained_records
+        ],
+        {
+            gene_key: isoforms
+            for gene_key, isoforms in reference_genes.items()
+            if gene_key[0] in retained_records
+        },
+    )
+
+
+def required_file(directory, filename):
+    """Return a required file path with a useful error when it is absent."""
+    path = Path(directory) / filename
+    if not path.is_file():
+        raise FileNotFoundError(f"required evaluation file not found: {path}")
+    return path
+
+
+def default_codon_table_path(genome_dir):
+    """Return an accession-specific codon table when one is bundled."""
+    codon_table_directory = Path(__file__).resolve().parent.parent / "codon_tables"
+    accession_path = codon_table_directory / f"{Path(genome_dir).name}.json"
+    if accession_path.is_file():
+        return accession_path
+    return codon_table_directory / "standard.json"
+
+
+def evaluate_prediction(
+    prediction_prefix,
+    genome_dir,
+    confidence_threshold=CONFIDENCE_THRESHOLD,
+    codon_table_path=None,
+):
+    """Load one prediction/reference pair and return count and rate metrics.
+
+    Gene and protein metrics count the predicted genes whose confidence is at
+    least ``confidence_threshold``; genes without a confidence
+    (``confidence=NA``) are always counted, and ``0`` counts every predicted
+    gene. Exon, boundary, and nucleotide metrics use every predicted CDS
+    interval. Without ``codon_table_path`` the translation table is
+    ``codon_tables/<accession>.json`` when the genome directory is named after
+    an accession with a bundled table, otherwise the standard table.
+    """
+    prediction_prefix = Path(prediction_prefix)
+    genome_dir = Path(genome_dir)
+    if (
+        not math.isfinite(confidence_threshold)
+        or not 0.0 <= confidence_threshold <= 1.0
+    ):
+        raise ValueError("confidence_threshold must be between 0 and 1")
+    codon_table_path = (
+        Path(codon_table_path)
+        if codon_table_path is not None
+        else default_codon_table_path(genome_dir)
+    )
+    codon_to_amino_acid = load_codon_table(codon_table_path)
+    retained_records = load_nuclear_evaluation_accessions(genome_dir)
+
+    predicted_gff_path = str(prediction_prefix) + ".gff"
+    predicted_gene_scores = parse_predicted_gene_scores(predicted_gff_path)
+    predicted_tracks = parse_predicted_cds_tracks(predicted_gff_path)
+    predicted_cds_frame_records = [
+        prediction
+        for prediction in parse_predicted_cds_records_with_frames(
+            str(prediction_prefix) + ".fna"
+        )
+        if confidence_threshold <= 0.0
+        or predicted_gene_scores.get(prediction[0]) is None
+        or predicted_gene_scores[prediction[0]] >= confidence_threshold
+    ]
+    reference_tracks = parse_reference_cds_tracks(
+        required_file(genome_dir, "genomic.gff"),
+        retained_records,
+    )
+    reference_cds_path = required_file(genome_dir, "cds_from_genomic.fna")
+    reference_genes = parse_reference_cds_fasta(reference_cds_path)
+    reference_proteins = parse_reference_translated_proteins(
+        reference_cds_path,
+        codon_to_amino_acid,
+    )
+
+    predicted_tracks = filter_tracks(predicted_tracks, retained_records)
+    reference_tracks = filter_tracks(reference_tracks, retained_records)
+    if retained_records is not None:
+        predicted_cds_frame_records = [
+            prediction
+            for prediction in predicted_cds_frame_records
+            if prediction[1] in retained_records
+        ]
+    predicted_cds = [
+        (record_id, sequence)
+        for _, record_id, sequence, _ in predicted_cds_frame_records
+    ]
+    predicted_proteins = [
+        (record_id, translate_cds_for_evaluation(sequence, frame, codon_to_amino_acid))
+        for _, record_id, sequence, frame in predicted_cds_frame_records
+    ]
+    _, reference_genes = filter_gene_inputs([], reference_genes, retained_records)
+    _, reference_proteins = filter_gene_inputs([], reference_proteins, retained_records)
+
+    matched_gene_by_prediction = gene_exact_match_assignments(
+        predicted_cds,
+        reference_genes,
+    )
+    metrics = gene_metric_counts(
+        predicted_cds,
+        reference_genes,
+        matched_gene_by_prediction,
+    )
+    metrics.update(protein_metric_counts(
+        predicted_proteins,
+        reference_proteins,
+        matched_gene_by_prediction,
+    ))
+    metrics.update(confidence_exact_match_metrics(
+        [
+            predicted_gene_scores.get(gene_id)
+            for gene_id, _, _, _ in predicted_cds_frame_records
+        ],
+        matched_gene_by_prediction,
+    ))
+    metrics.update(interval_metric_counts(predicted_tracks, reference_tracks))
+    add_metric_rates(metrics)
+
+    metrics["genome"] = genome_dir.name
+    metrics["codon_table"] = codon_table_path.name
+    metrics["confidence_threshold"] = confidence_threshold
+    metrics["records"] = len(
+        {record_id for record_id, _ in predicted_tracks}
+        | {record_id for record_id, _ in reference_tracks}
+    )
+    metrics["record_filter"] = (
+        "all (no sequence report)" if retained_records is None else "nuclear"
+    )
+    return metrics
+
+
+def write_evaluation_metrics_json(path, metrics):
+    """Atomically write the evaluation metrics of one accession."""
+    output_path = Path(path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "accession": metrics["genome"],
+        "generated_at": datetime.now(timezone.utc).isoformat(),
+        "metrics": metrics,
+    }
+    temporary_path = output_path.with_name(output_path.name + ".tmp")
+    with temporary_path.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, sort_keys=True, allow_nan=False)
+        handle.write("\n")
+    temporary_path.replace(output_path)
+
+
+def print_evaluation_metrics(metrics):
+    """Print a concise summary of one evaluation."""
+    print(
+        f"[evaluation] {metrics['genome']} filter={metrics['record_filter']} "
+        f"codon_table={metrics['codon_table']} "
+        f"confidence_threshold={metrics['confidence_threshold']:.2f} "
+        f"({metrics['records']} records, "
+        f"{metrics['reference_genes']} reference genes, "
+        f"{metrics['predicted_genes']} predictions)"
+    )
+    for label, metric_name in (
+        ("gene", "gene"),
+        ("protein", "protein"),
+        ("exon", "exon"),
+        ("5' boundary", "five_prime_boundary"),
+        ("3' boundary", "three_prime_boundary"),
+        ("nucleotide", "nucleotide"),
+    ):
+        print(
+            f"  {label:<11} P/R/F1: "
+            f"{metrics[f'{metric_name}_precision']:.4f} / "
+            f"{metrics[f'{metric_name}_recall']:.4f} / "
+            f"{metrics[f'{metric_name}_f1']:.4f}"
+        )
+    scored_genes = metrics.get("confidence_scored_genes")
+    if scored_genes:
+        correlation = metrics["confidence_exact_match_correlation"]
+        correlation_text = "NA" if correlation is None else f"{correlation:.4f}"
+        auroc = metrics["confidence_exact_match_auroc"]
+        auprc = metrics["confidence_exact_match_auprc"]
+        auroc_text = "NA" if auroc is None else f"{auroc:.4f}"
+        auprc_text = "NA" if auprc is None else f"{auprc:.4f}"
+        exact_mean = metrics["confidence_exact_mean"]
+        nonexact_mean = metrics["confidence_nonexact_mean"]
+        exact_mean_text = "NA" if exact_mean is None else f"{exact_mean:.4f}"
+        nonexact_mean_text = (
+            "NA" if nonexact_mean is None else f"{nonexact_mean:.4f}"
+        )
+        print(
+            "  confidence vs gene exact match: "
+            f"AUROC/AUPRC={auroc_text}/{auprc_text}; "
+            f"point-biserial r={correlation_text}; "
+            f"mean exact/non-exact={exact_mean_text}/{nonexact_mean_text}; "
+            f"exact/non-exact={metrics['confidence_exact_genes']}/"
+            f"{metrics['confidence_nonexact_genes']}; "
+            f"scored/missing={scored_genes}/"
+            f"{metrics['confidence_missing_genes']}"
+        )
+    elif scored_genes == 0:
+        print("  confidence vs gene exact match: unavailable (no scores)")
+
+
+def parse_args():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "--pred",
+        required=True,
+        help="prediction prefix; reads <prefix>.gff and <prefix>.fna",
+    )
+    parser.add_argument(
+        "--genome_dir",
+        required=True,
+        help=(
+            "NCBI directory containing genomic.gff and cds_from_genomic.fna; "
+            "optional sequence_report.jsonl enables nuclear-record selection"
+        ),
+    )
+    parser.add_argument(
+        "--codon-table",
+        default=None,
+        help=(
+            "JSON translation table for predicted and reference CDS; defaults "
+            "to codon_tables/<accession>.json when present, otherwise "
+            "codon_tables/standard.json"
+        ),
+    )
+    parser.add_argument(
+        "--confidence-threshold",
+        type=float,
+        default=CONFIDENCE_THRESHOLD,
+        help=(
+            "gene-level confidence threshold: gene and protein metrics count "
+            "predicted genes with confidence at or above this value "
+            f"(default {CONFIDENCE_THRESHOLD}); genes without a confidence are "
+            "always counted and 0 counts every gene"
+        ),
+    )
+    parser.add_argument(
+        "--output-json",
+        help="atomically write the metrics to this JSON file",
+    )
+    return parser.parse_args()
+
+
+def main():
+    """Command-line entry point."""
+    args = parse_args()
+    if (
+        not math.isfinite(args.confidence_threshold)
+        or not 0.0 <= args.confidence_threshold <= 1.0
+    ):
+        raise SystemExit("--confidence-threshold must be between 0 and 1")
+
+    metrics = evaluate_prediction(
+        args.pred,
+        args.genome_dir,
+        confidence_threshold=args.confidence_threshold,
+        codon_table_path=args.codon_table,
+    )
+    print_evaluation_metrics(metrics)
+    if args.output_json:
+        write_evaluation_metrics_json(args.output_json, metrics)
+        print(f"[evaluation] wrote {args.output_json}")
+    return metrics
+
+
+if __name__ == "__main__":
+    main()

annotation_pipeline/src/model_inference.py

@@ -0,0 +1,1067 @@
+import os
+import argparse
+import datetime
+import gzip
+import signal
+import threading
+import queue
+import time
+from typing import Any, Dict, IO, List, Optional, Sequence, Tuple
+
+import numpy as np
+import pandas as pd
+
+try:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    HAS_PYARROW = True
+except ImportError:
+    pa = None
+    pq = None
+    HAS_PYARROW = False
+import torch
+import torch.multiprocessing as mp
+from tqdm import tqdm
+from transformers import (
+    AutoModelForTokenClassification,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+from constants import (
+    CONTEXT_LENGTH,
+    HEAD_NAMES,
+    MODEL_NAME,
+    OVERLAP_LENGTH,
+)
+
+FASTA_SUFFIXES = (".fasta", ".fa", ".fna", ".ffn", ".faa", ".fas")
+GENBANK_SUFFIXES = (".gb", ".gbk", ".gbff")
+
+
+def parse_arguments() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Pure model inference for CDS annotation. Outputs raw softmax probabilities to parquet.",
+    )
+
+    parser.add_argument(
+        "--input",
+        type=str,
+        nargs="+",
+        required=True,
+        help="Genome FASTA or GenBank files (optionally gzip-compressed) or directories",
+    )
+    parser.add_argument(
+        "--output_path", type=str, default="./annotation_results", help="Output directory"
+    )
+
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        default=MODEL_NAME,
+        help="Hugging Face model id or local model directory",
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default=None,
+        help=(
+            "Model revision on the Hugging Face Hub (branch, tag, or commit); it "
+            "applies to the weights and the model code. Default: the latest revision."
+        ),
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=8, help="Batch size for inference"
+    )
+    parser.add_argument(
+        "--dispatch_chunk_size",
+        type=int,
+        default=4,
+        help=(
+            "Number of sequences per dynamic scheduler task. Smaller values improve "
+            "load balancing/work stealing; larger values improve batching locality."
+        ),
+    )
+    parser.add_argument(
+        "--gpu_count", type=int, default=-1, help="Number of GPUs to use (-1 for all)"
+    )
+    parser.add_argument(
+        "--attn_implementation",
+        type=str,
+        choices=["eager", "sdpa", "flash_attention_2"],
+        default="flash_attention_2",
+        help=(
+            "Attention implementation. flash_attention_2 (default) was used for all "
+            "reported results; sdpa needs no flash-attn package, but its "
+            "probabilities are not bit-identical."
+        ),
+    )
+
+    parser.add_argument(
+        "--context_length",
+        type=int,
+        default=CONTEXT_LENGTH,
+        help="Context length in tokens",
+    )
+    parser.add_argument(
+        "--overlap_length",
+        type=int,
+        default=OVERLAP_LENGTH,
+        help="Overlap length in tokens",
+    )
+
+    parser.add_argument(
+        "--prob_dtype",
+        type=str,
+        choices=["float16", "float32"],
+        default="float16",
+        help="Data type for saved CDS probabilities in parquet",
+    )
+    parser.add_argument(
+        "--save_non_coding_prob",
+        action="store_true",
+        help=(
+            "Also save non-coding probability. By default only P(CDS) is stored, which is lossless for binary softmax "
+            "because P(non_coding)=1-P(CDS)."
+        ),
+    )
+    parser.add_argument(
+        "--limit", type=int, default=None, help="Limit number of sequences (debug)"
+    )
+
+    return parser.parse_args()
+
+
+def open_text(path: str, errors: str = "strict") -> IO[str]:
+    """Open a text file, decompressing it when the name ends with .gz."""
+    if path.lower().endswith(".gz"):
+        return gzip.open(path, "rt", encoding="utf-8", errors=errors)
+    return open(path, "r", encoding="utf-8", errors=errors)
+
+
+def strip_gzip_suffix(name: str) -> str:
+    return name[: -len(".gz")] if name.lower().endswith(".gz") else name
+
+
+def parse_fasta_from_stream(stream: IO[str]) -> List[Tuple[str, str]]:
+    records: List[Tuple[str, str]] = []
+    header: Optional[str] = None
+    seq_lines: List[str] = []
+    for line in stream:
+        line = line.strip()
+        if line.startswith(">"):
+            if header is not None:
+                records.append((header, "".join(seq_lines).upper()))
+            header, seq_lines = line, []
+        else:
+            seq_lines.append(line)
+    if header is not None:
+        records.append((header, "".join(seq_lines).upper()))
+    return records
+
+
+def read_fasta(path: str) -> List[Tuple[str, str]]:
+    with open_text(path) as f:
+        records = parse_fasta_from_stream(f)
+    print(f"✅ Read {len(records)} sequences from local file: {path}")
+    return records
+
+
+def read_genbank(path: str) -> List[Tuple[str, str]]:
+    """Read the sequence of every GenBank record; the record name is its accession.version."""
+    from Bio import SeqIO
+    from Bio.Seq import UndefinedSequenceError
+
+    records: List[Tuple[str, str]] = []
+    with open_text(path) as f:
+        for record in SeqIO.parse(f, "genbank"):
+            try:
+                sequence = str(record.seq).upper()
+            except UndefinedSequenceError:
+                raise ValueError(
+                    f"GenBank record {record.id} in '{path}' contains no sequence"
+                ) from None
+            records.append((f">{record.id}", sequence))
+    print(f"✅ Read {len(records)} sequences from GenBank file: {path}")
+    return records
+
+
+def sniff_input_format(path: str, lines_to_check: int = 50) -> Optional[str]:
+    with open_text(path, errors="ignore") as f:
+        for _ in range(lines_to_check):
+            line = f.readline()
+            if not line:
+                break
+            line = line.strip()
+            if line:
+                if line.startswith(">"):
+                    return "fasta"
+                if line.startswith("LOCUS"):
+                    return "genbank"
+                return None
+    return None
+
+
+def detect_input_format(path: str) -> str:
+    name = strip_gzip_suffix(path.lower())
+    if name.endswith(FASTA_SUFFIXES):
+        return "fasta"
+    if name.endswith(GENBANK_SUFFIXES):
+        return "genbank"
+    input_format = sniff_input_format(path)
+    if input_format is None:
+        raise ValueError(
+            f"Unsupported input format for '{path}'. Please provide a FASTA or GenBank file."
+        )
+    return input_format
+
+
+def resolve_input_paths(input_items: Sequence[str]) -> List[str]:
+    supported_suffixes = FASTA_SUFFIXES + GENBANK_SUFFIXES
+    resolved: List[str] = []
+    for item in input_items:
+        if os.path.isdir(item):
+            files_in_dir: List[str] = []
+            for root, _, files in os.walk(item):
+                for name in files:
+                    if strip_gzip_suffix(name.lower()).endswith(supported_suffixes):
+                        files_in_dir.append(os.path.join(root, name))
+            resolved.extend(sorted(files_in_dir))
+            continue
+
+        resolved.append(item)
+    return resolved
+
+
+def get_record_name(header: str) -> str:
+    text = header[1:] if header.startswith(">") else header
+    text = text.strip()
+    return text.split()[0] if text else "record"
+
+
+def pyarrow_list_array_from_ndarrays(
+    arrays: Sequence[np.ndarray], value_type: "pa.DataType"
+) -> "pa.Array":
+    assert HAS_PYARROW and pa is not None
+    return pa.array(
+        [pa.array(np.asarray(arr), type=value_type) for arr in arrays],
+        type=pa.list_(value_type),
+    )
+
+
+def write_prob_parquet(
+    records: Sequence[Tuple[str, str]],
+    probs_per_head: Sequence[Sequence[np.ndarray]],
+    head_names: Sequence[str],
+    output_path: str,
+    prob_dtype: str,
+    save_non_coding_prob: bool,
+) -> None:
+    if not probs_per_head:
+        raise ValueError("No probabilities to write.")
+
+    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+    np_dtype = np.float16 if prob_dtype == "float16" else np.float32
+
+    record_names = [get_record_name(h) for h, _ in records]
+    sequences = [seq for _, seq in records]
+
+    if HAS_PYARROW and pa is not None and pq is not None:
+        value_type = pa.float16() if prob_dtype == "float16" else pa.float32()
+        cols: Dict[str, Any] = {}
+        cols["record_name"] = pa.array(record_names, type=pa.string())
+        cols["sequence"] = pa.array(sequences, type=pa.string())
+
+        for h, head_name in enumerate(head_names):
+            cds_arrays = [np.asarray(arr, dtype=np_dtype) for arr in probs_per_head[h]]
+            cols[f"pred_prob_{head_name}_cds"] = pyarrow_list_array_from_ndarrays(
+                cds_arrays, value_type
+            )
+            if save_non_coding_prob:
+                nc_arrays = [
+                    np.asarray(1.0 - arr, dtype=np_dtype) for arr in cds_arrays
+                ]
+                cols[f"pred_prob_{head_name}_non_coding"] = (
+                    pyarrow_list_array_from_ndarrays(nc_arrays, value_type)
+                )
+
+        table = pa.table(cols)
+        pq.write_table(table, output_path, compression="zstd")
+    else:
+        data: List[Dict[str, Any]] = []
+        for i in range(len(records)):
+            row: Dict[str, Any] = {
+                "record_name": record_names[i],
+                "sequence": sequences[i],
+            }
+            for h, head_name in enumerate(head_names):
+                row[f"pred_prob_{head_name}_cds"] = np.asarray(
+                    probs_per_head[h][i], dtype=np_dtype
+                ).tolist()
+                if save_non_coding_prob:
+                    row[f"pred_prob_{head_name}_non_coding"] = np.asarray(
+                        1.0 - probs_per_head[h][i], dtype=np_dtype
+                    ).tolist()
+            data.append(row)
+        df = pd.DataFrame(data)
+        df.to_parquet(output_path, index=False)
+
+    print(f"✅ Raw probability parquet written to {output_path}")
+
+
+def setup_model_for_gpu(
+    model_name: str,
+    gpu_id: int,
+    dtype_str: str,
+    attn_implementation: Optional[str] = None,
+    revision: Optional[str] = None,
+) -> Tuple[PreTrainedModel, PreTrainedTokenizer, torch.device]:
+    revision_note = f" (revision {revision})" if revision else ""
+    print(f"🤗 Loading model on GPU {gpu_id}: {model_name}{revision_note}")
+    start_time = time.time()
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name, revision=revision, trust_remote_code=True
+    )
+    model_kwargs: Dict[str, Any] = {
+        "dtype": getattr(torch, dtype_str),
+        "revision": revision,
+        "trust_remote_code": True,
+    }
+    if attn_implementation is not None:
+        model_kwargs["attn_implementation"] = attn_implementation
+
+    model = AutoModelForTokenClassification.from_pretrained(
+        model_name,
+        **model_kwargs,
+
+    )
+
+    device = torch.device(f"cuda:{gpu_id}" if gpu_id >= 0 else "cpu")
+    model.to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    effective_attn = getattr(model.config, "_attn_implementation", None)
+    first_attn_class = None
+    base_model = getattr(model, "model", None)
+    layers = getattr(base_model, "layers", None)
+    if layers:
+        first_attn = getattr(layers[0], "self_attn", None)
+        if first_attn is not None:
+            first_attn_class = type(first_attn).__name__
+    print(f"📊 Device {device} model size: {total_params / 1e6:.1f}M parameters")
+    print(
+        f"🧠 Attention implementation: requested={attn_implementation or 'default'} "
+        f"effective={effective_attn} first_layer={first_attn_class}"
+    )
+    print(f"⏱️ Model loading completed in {time.time() - start_time:.2f} seconds")
+    return model, tokenizer, device
+
+
+def distribute_sequences_to_gpus(
+    records: Sequence[Tuple[str, str]],
+    gpu_count: int,
+    input_ids_per_record: Optional[Sequence[np.ndarray]] = None,
+) -> Tuple[List[List[Tuple[str, str]]], Optional[List[List[np.ndarray]]]]:
+    total = len(records)
+    # Use contiguous, workload-weighted partitioning (by sequence length) so that
+    # global record order is preserved while reducing multi-GPU stragglers.
+    weights = [max(1, len(seq)) for _, seq in records]
+    prefix = [0]
+    for w in weights:
+        prefix.append(prefix[-1] + w)
+
+    def range_weight(start: int, end: int) -> int:
+        return prefix[end] - prefix[start]
+
+    gpu_sequences: List[List[Tuple[str, str]]] = []
+    gpu_input_ids: Optional[List[List[np.ndarray]]] = (
+        [] if input_ids_per_record is not None else None
+    )
+
+    cursor = 0
+    for gpu_id in range(gpu_count):
+        remaining_gpus = gpu_count - gpu_id
+        remaining_records = total - cursor
+
+        if remaining_records <= 0:
+            end = cursor
+        elif remaining_gpus == 1:
+            end = total
+        else:
+            max_end = total - (remaining_gpus - 1)  # leave >=1 record per later GPU
+            remaining_weight = range_weight(cursor, total)
+            target = remaining_weight / remaining_gpus
+
+            end = cursor
+            accum = 0
+            while end < max_end:
+                next_w = weights[end]
+                if end > cursor and accum >= target:
+                    break
+                accum += next_w
+                end += 1
+
+            if end > cursor + 1:
+                with_last = abs(accum - target)
+                without_last = abs((accum - weights[end - 1]) - target)
+                if without_last <= with_last:
+                    end -= 1
+
+        chunk = list(records[cursor:end])
+        gpu_sequences.append(chunk)
+        if gpu_input_ids is not None:
+            gpu_input_ids.append(list(input_ids_per_record[cursor:end]))
+
+        assigned_bp = range_weight(cursor, end)
+        print(
+            f"📋 GPU {gpu_id} assigned {len(chunk)} sequences (~{assigned_bp} bp)"
+        )
+        cursor = end
+
+    return gpu_sequences, gpu_input_ids
+
+
+@torch.no_grad()
+def process_sequences_on_gpu(
+    sequences: Sequence[Tuple[str, str]],
+    model: PreTrainedModel,
+    tokenizer: PreTrainedTokenizer,
+    device: torch.device,
+    max_length: int,
+    overlap_length: int,
+    micro_batch_size: int,
+    progress_event_queue: mp.Queue,
+) -> List[List[np.ndarray]]:
+    pad_id = tokenizer.pad_token_id
+    num_heads = getattr(model, "num_prediction_heads", 1)
+    tokenizer_k = getattr(tokenizer, "k", 1) or 1
+
+    max_char_length = max_length * tokenizer_k
+    overlap_char_length = overlap_length * tokenizer_k
+    num_sequences = len(sequences)
+
+    prob_per_head: List[List[np.ndarray]] = [
+        [np.empty((0,), dtype=np.float32) for _ in range(num_sequences)]
+        for _ in range(num_heads)
+    ]
+
+    ChunkTask = Tuple[int, int, int, List[int], List[int]]
+    seq_cds_sums: List[Optional[List[torch.Tensor]]] = [None] * num_sequences
+    seq_overlap_cnt: List[Optional[torch.Tensor]] = [None] * num_sequences
+    remaining_chunks = [0] * num_sequences
+    full_pending: List[ChunkTask] = []
+    padded_pending: List[ChunkTask] = []
+
+    def windows_for_sequence(seq_len: int) -> List[Tuple[int, int]]:
+        windows: List[Tuple[int, int]] = []
+        if seq_len < max_char_length:
+            main_len = (seq_len // tokenizer_k) * tokenizer_k
+            if main_len:
+                windows.append((0, main_len))
+            tail_offset = seq_len % tokenizer_k
+            if tail_offset:
+                windows.append((tail_offset, seq_len - tail_offset))
+        else:
+            seq_start = 0
+            seq_work_len = seq_len
+            while True:
+                char_length = min(max_char_length, seq_work_len)
+                windows.append((seq_start, char_length))
+                if char_length == seq_work_len:
+                    break
+                if (
+                    seq_work_len - max_char_length + overlap_char_length
+                    < max_char_length
+                ):
+                    seq_start = seq_len - max_char_length
+                    seq_work_len = max_char_length
+                else:
+                    step_chars = max_char_length - overlap_char_length
+                    seq_start += step_chars
+                    seq_work_len -= step_chars
+        return windows
+
+    def finalize_sequence(seq_idx: int) -> None:
+        counts = seq_overlap_cnt[seq_idx]
+        sums = seq_cds_sums[seq_idx]
+        if counts is None or sums is None:
+            raise RuntimeError(f"Missing accumulation buffers for sequence {seq_idx}")
+        if torch.any(counts == 0):
+            raise RuntimeError(f"Sequence {seq_idx} has uncovered base positions")
+
+        denom = counts.to(dtype=torch.float32)
+        for h in range(num_heads):
+            prob_per_head[h][seq_idx] = (
+                (sums[h] / denom).numpy().astype(np.float32, copy=False)
+            )
+        seq_overlap_cnt[seq_idx] = None
+        seq_cds_sums[seq_idx] = None
+        progress_event_queue.put("infer")
+
+    def run_batch(batch: Sequence[ChunkTask]) -> None:
+        if not batch:
+            return
+        inp = torch.tensor([task[3] for task in batch], dtype=torch.long).to(device)
+        att = torch.tensor([task[4] for task in batch], dtype=torch.long).to(device)
+
+        # Count total tokens in this batch (sum of attention mask)
+        total_tokens = att.sum().item()
+
+        # Full-window batches carry all-ones masks. GENERanno's token-
+        # classification head needs the mask for output sizing, and the base
+        # model drops all-ones masks internally so FA2 can avoid padding-mask
+        # fallback. Padded windows are kept in separate masked batches.
+        logits = model(input_ids=inp, attention_mask=att).logits
+        probs = logits.softmax(dim=-1).cpu()
+
+        for local_i, (seq_idx, char_start, char_length, _, mask) in enumerate(batch):
+            valid_len = int(sum(mask)) * tokenizer_k
+            if valid_len != char_length:
+                raise RuntimeError(
+                    f"Chunk length mismatch: valid_len={valid_len}, "
+                    f"char_length={char_length}"
+                )
+            total_pred_len = valid_len * num_heads
+            chunk_probs_all_heads = probs[local_i, :total_pred_len]
+            chunk_probs_all_heads = chunk_probs_all_heads.view(
+                num_heads, valid_len, -1
+            )
+
+            counts = seq_overlap_cnt[seq_idx]
+            sums = seq_cds_sums[seq_idx]
+            if counts is None or sums is None:
+                raise RuntimeError(
+                    f"Missing accumulation buffers for sequence {seq_idx}"
+                )
+            counts[char_start : char_start + char_length] += 1
+            for h in range(num_heads):
+                sums[h][char_start : char_start + char_length] += (
+                    chunk_probs_all_heads[h, :, 1].float()
+                )
+
+            remaining_chunks[seq_idx] -= 1
+            if remaining_chunks[seq_idx] == 0:
+                finalize_sequence(seq_idx)
+
+    def flush_ready(pending: List[ChunkTask]) -> None:
+        while len(pending) >= micro_batch_size:
+            batch = pending[:micro_batch_size]
+            del pending[:micro_batch_size]
+            run_batch(batch)
+
+    def flush_all(pending: List[ChunkTask]) -> None:
+        while pending:
+            batch = pending[:micro_batch_size]
+            del pending[:micro_batch_size]
+            run_batch(batch)
+
+    for seq_idx, (_, seq) in enumerate(sequences):
+        if len(seq) < tokenizer_k:
+            for h in range(num_heads):
+                prob_per_head[h][seq_idx] = np.zeros((len(seq),), dtype=np.float32)
+            progress_event_queue.put("infer")
+            continue
+
+        windows = windows_for_sequence(len(seq))
+        if not windows:
+            for h in range(num_heads):
+                prob_per_head[h][seq_idx] = np.zeros((len(seq),), dtype=np.float32)
+            progress_event_queue.put("infer")
+            continue
+
+        remaining_chunks[seq_idx] = len(windows)
+        seq_overlap_cnt[seq_idx] = torch.zeros(len(seq), dtype=torch.long)
+        seq_cds_sums[seq_idx] = [
+            torch.zeros(len(seq), dtype=torch.float32) for _ in range(num_heads)
+        ]
+
+        for char_start, char_length in windows:
+            chrs = seq[char_start : char_start + char_length]
+            chunk_seq = tokenizer(chrs, add_special_tokens=False)["input_ids"]
+
+            if len(chunk_seq) > max_length:
+                raise ValueError(
+                    f"Chunk has {len(chunk_seq)} tokens, exceeding max_length={max_length}"
+                )
+            valid_token_len = len(chunk_seq)
+            if valid_token_len * tokenizer_k != char_length:
+                raise RuntimeError(
+                    f"Tokenized chunk length mismatch: {valid_token_len} tokens "
+                    f"for {char_length} bp"
+                )
+            pad_len = max_length - valid_token_len
+            mask = [1] * valid_token_len + [0] * pad_len
+            task: ChunkTask = (
+                seq_idx,
+                char_start,
+                char_length,
+                chunk_seq + [pad_id] * pad_len,
+                mask,
+            )
+            if valid_token_len == max_length:
+                full_pending.append(task)
+                flush_ready(full_pending)
+            else:
+                padded_pending.append(task)
+                flush_ready(padded_pending)
+
+    flush_all(full_pending)
+    flush_all(padded_pending)
+
+    return prob_per_head
+
+
+def persistent_worker_process(
+    gpu_id: int,
+    model_name: str,
+    revision: Optional[str],
+    dtype_str: str,
+    attn_implementation: Optional[str],
+    max_length: int,
+    overlap_length: int,
+    micro_batch_size: int,
+    task_queue: mp.Queue,
+    result_queue: mp.Queue,
+    progress_event_queue: mp.Queue,
+) -> None:
+    job_id = -1
+    task_id = -1
+    try:
+        model: Optional[PreTrainedModel] = None
+        tokenizer: Optional[PreTrainedTokenizer] = None
+        device: Optional[torch.device] = None
+
+        while True:
+            task_item = task_queue.get()
+            if task_item is None:
+                break
+
+            job_id, task_id, task_seq_indices, sequences = task_item
+            if model is None:
+                model, tokenizer, device = setup_model_for_gpu(
+                    model_name, gpu_id, dtype_str, attn_implementation, revision
+                )
+            assert tokenizer is not None and device is not None and model is not None
+            gpu_results = process_sequences_on_gpu(
+                sequences=sequences,
+                model=model,
+                tokenizer=tokenizer,
+                device=device,
+                max_length=max_length,
+                overlap_length=overlap_length,
+                micro_batch_size=micro_batch_size,
+                progress_event_queue=progress_event_queue,
+            )
+            result_queue.put(
+                ("ok", job_id, gpu_id, task_id, task_seq_indices, gpu_results)
+            )
+    except KeyboardInterrupt:
+        pass
+    except Exception as e:
+        print(f"❌ Error in persistent GPU {gpu_id} worker: {e}")
+        result_queue.put(("error", job_id, gpu_id, task_id, str(e)))
+
+
+def create_persistent_worker_runtime(
+    model_name: str,
+    revision: Optional[str],
+    dtype_str: str,
+    attn_implementation: Optional[str],
+    gpu_count: int,
+    max_length: int,
+    overlap_length: int,
+    micro_batch_size: int,
+) -> Dict[str, Any]:
+    task_queue = mp.Queue()
+    result_queue = mp.Queue()
+    progress_event_queue = mp.Queue()
+
+    processes: List[mp.Process] = []
+    for gpu_id in range(gpu_count):
+        p = mp.Process(
+            target=persistent_worker_process,
+            args=(
+                gpu_id,
+                model_name,
+                revision,
+                dtype_str,
+                attn_implementation,
+                max_length,
+                overlap_length,
+                micro_batch_size,
+                task_queue,
+                result_queue,
+                progress_event_queue,
+            ),
+        )
+        p.start()
+        processes.append(p)
+
+    return {
+        "gpu_count": gpu_count,
+        "task_queue": task_queue,
+        "result_queue": result_queue,
+        "progress_event_queue": progress_event_queue,
+        "processes": processes,
+        "next_job_id": 0,
+    }
+
+
+def shutdown_persistent_worker_runtime(
+    runtime: Dict[str, Any], interrupted: bool = False
+) -> None:
+    task_queue = runtime["task_queue"]
+    result_queue = runtime["result_queue"]
+    progress_event_queue = runtime["progress_event_queue"]
+    processes = runtime["processes"]
+
+    if interrupted:
+        for p in processes:
+            if p.is_alive() and p.pid is not None:
+                os.kill(p.pid, signal.SIGINT)
+    else:
+        for _ in processes:
+            task_queue.put(None)
+
+    deadline = time.monotonic() + 30
+    for p in processes:
+        p.join(timeout=max(0, deadline - time.monotonic()))
+    for p in processes:
+        if p.is_alive():
+            p.terminate()
+            p.join()
+
+    task_queue.close()
+    # Tasks left behind by a worker that failed (for example to load the
+    # model) can no longer be read; waiting for them to be sent would hang.
+    task_queue.cancel_join_thread()
+    result_queue.close()
+    result_queue.join_thread()
+    progress_event_queue.close()
+    progress_event_queue.join_thread()
+
+
+def build_dynamic_sequence_tasks(
+    records: Sequence[Tuple[str, str]],
+    dispatch_chunk_size: int,
+) -> List[Tuple[int, List[int], List[Tuple[str, str]]]]:
+    if dispatch_chunk_size <= 0:
+        raise ValueError(
+            f"--dispatch_chunk_size must be a positive integer, got {dispatch_chunk_size}"
+        )
+
+    # Longest-first scheduling reduces end-of-job stragglers when sequence sizes vary.
+    ordered_indices = sorted(range(len(records)), key=lambda i: len(records[i][1]), reverse=True)
+    tasks: List[Tuple[int, List[int], List[Tuple[str, str]]]] = []
+    for task_id, start in enumerate(range(0, len(ordered_indices), dispatch_chunk_size)):
+        seq_indices = ordered_indices[start : start + dispatch_chunk_size]
+        seq_chunk = [records[i] for i in seq_indices]
+        tasks.append((task_id, seq_indices, seq_chunk))
+    return tasks
+
+
+def infer_records(
+    records: Sequence[Tuple[str, str]],
+    model_name: str,
+    dtype_str: str,
+    gpu_count: int,
+    max_length: int,
+    overlap_length: int,
+    micro_batch_size: int,
+    dispatch_chunk_size: int,
+    persistent_runtime: Optional[Dict[str, Any]] = None,
+    attn_implementation: Optional[str] = None,
+    revision: Optional[str] = None,
+) -> List[List[np.ndarray]]:
+    owns_runtime = False
+    runtime = persistent_runtime
+    if runtime is None:
+        print(f"🚀 Starting persistent GPU workers ({gpu_count} GPUs)")
+        runtime = create_persistent_worker_runtime(
+            model_name=model_name,
+            revision=revision,
+            dtype_str=dtype_str,
+            attn_implementation=attn_implementation,
+            gpu_count=gpu_count,
+            max_length=max_length,
+            overlap_length=overlap_length,
+            micro_batch_size=micro_batch_size,
+        )
+        owns_runtime = True
+    else:
+        print(f"🚀 Reusing persistent GPU workers ({gpu_count} GPUs)")
+
+    if int(runtime["gpu_count"]) != gpu_count:
+        raise ValueError(
+            f"persistent runtime gpu_count={runtime['gpu_count']} does not match requested gpu_count={gpu_count}"
+        )
+
+    interrupted = False
+    try:
+        tasks = build_dynamic_sequence_tasks(
+            records=records,
+            dispatch_chunk_size=dispatch_chunk_size,
+        )
+        print(
+            f"📊 Dynamic scheduling {len(records)} sequences across {gpu_count} GPUs "
+            f"using {len(tasks)} task(s) (chunk_size={dispatch_chunk_size})"
+        )
+
+        result_queue = runtime["result_queue"]
+        progress_event_queue = runtime["progress_event_queue"]
+        task_queue = runtime["task_queue"]
+
+        job_id = runtime["next_job_id"]
+        runtime["next_job_id"] = job_id + 1
+        for task_id, seq_indices, seq_chunk in tasks:
+            task_queue.put((job_id, task_id, seq_indices, seq_chunk))
+
+        total_sequences = len(records)
+        num_heads = 0
+        all_probs_per_head: List[List[Optional[np.ndarray]]] = []
+        expected_results = len(tasks)
+        received_results = 0
+
+        def progress_worker() -> None:
+            infer_pbar = tqdm(
+                total=total_sequences,
+                desc="Inference",
+                unit="seq",
+                position=0,
+                leave=True,
+                dynamic_ncols=True,
+                miniters=1,
+                mininterval=0.0,
+            )
+            while True:
+                event_type = progress_event_queue.get()
+                if event_type == "stop":
+                    break
+                if event_type == "infer":
+                    infer_pbar.update(1)
+            infer_pbar.close()
+
+        progress_thread = threading.Thread(target=progress_worker, daemon=True)
+        progress_thread.start()
+        try:
+            while received_results < expected_results:
+                try:
+                    result = result_queue.get(
+                        timeout=30
+                    )
+                except queue.Empty:
+                    dead_workers = [
+                        i
+                        for i, p in enumerate(runtime["processes"])
+                        if not p.is_alive()
+                    ]
+                    if dead_workers:
+                        raise RuntimeError(
+                            f"Worker process(es) on GPU {dead_workers} died unexpectedly, "
+                            f"exit codes: {[runtime['processes'][i].exitcode for i in dead_workers]}"
+                        )
+                    continue
+                state = result[0]
+                result_job_id = result[1]
+                result_gpu_id = result[2]
+                if state == "error":
+                    error_task_id = result[3]
+                    payload = result[4]
+                    raise RuntimeError(
+                        f"Worker error on GPU {result_gpu_id} (task {error_task_id}): {payload}"
+                    )
+                if result_job_id != job_id:
+                    continue
+                _, _, _, result_task_id, task_seq_indices, task_probs = result
+                if num_heads == 0:
+                    num_heads = len(task_probs)
+                    all_probs_per_head = [
+                        [None] * total_sequences for _ in range(num_heads)
+                    ]
+                if len(task_probs) != num_heads:
+                    raise RuntimeError(
+                        f"Inconsistent number of heads from task {result_task_id}: "
+                        f"expected {num_heads}, got {len(task_probs)}"
+                    )
+                for local_i, seq_idx in enumerate(task_seq_indices):
+                    for head_idx in range(num_heads):
+                        all_probs_per_head[head_idx][seq_idx] = task_probs[head_idx][
+                            local_i
+                        ]
+                received_results += 1
+        except KeyboardInterrupt:
+            interrupted = True
+            raise
+        finally:
+            progress_event_queue.put("stop")
+            progress_thread.join()
+
+        print("🔗 Combining results from all GPUs...")
+        finalized_probs: List[List[np.ndarray]] = []
+        for head_idx in range(num_heads):
+            head_probs = all_probs_per_head[head_idx]
+            if any(p is None for p in head_probs):
+                missing = sum(1 for p in head_probs if p is None)
+                raise RuntimeError(
+                    f"Missing probabilities for {missing} sequence(s) in head {head_idx}"
+                )
+            finalized_probs.append([p for p in head_probs if p is not None])
+
+        print(
+            f"✅ Successfully processed {len(records)} sequences across {gpu_count} GPUs"
+        )
+        return finalized_probs
+    finally:
+        if owns_runtime:
+            shutdown_persistent_worker_runtime(runtime, interrupted=interrupted)
+
+
+def display_progress_header() -> None:
+    print("\n" + "=" * 80)
+    print("🧬  CDS RAW MODEL INFERENCE PIPELINE  🧬")
+    print("=" * 80 + "\n")
+
+
+def read_input_records(
+    input_path: str,
+    limit: Optional[int],
+) -> List[Tuple[str, str]]:
+    print("🔄 Reading input file...")
+    if detect_input_format(input_path) == "genbank":
+        records = read_genbank(input_path)
+    else:
+        records = read_fasta(input_path)
+    if limit is not None:
+        records = records[:limit]
+    return records
+
+
+def make_output_parquet_path(
+    output_dir: str, input_path: str, run_timestamp: str
+) -> str:
+    input_filename = strip_gzip_suffix(os.path.basename(input_path))
+    base_input_name = os.path.splitext(input_filename)[0]
+    return os.path.join(
+        output_dir, f"{base_input_name}_{run_timestamp}.raw_probs.parquet"
+    )
+
+
+def main() -> None:
+    display_progress_header()
+    total_start_time = time.time()
+    args = parse_arguments()
+
+    # Every reported result was computed with bfloat16 weights, which
+    # FlashAttention-2 also requires.
+    dtype_str = "bfloat16"
+    print(f"📊 Using model dtype: {dtype_str}")
+    print(f"💾 Saving probability dtype: {args.prob_dtype}")
+    if not args.save_non_coding_prob:
+        print(
+            "ℹ️ Saving only P(CDS) per head; P(non_coding)=1-P(CDS) for binary softmax."
+        )
+
+    available_gpus = torch.cuda.device_count()
+    if available_gpus <= 0:
+        raise RuntimeError("No CUDA devices available.")
+    if args.gpu_count == -1:
+        gpu_count = available_gpus
+    elif args.gpu_count <= 0:
+        raise ValueError(
+            f"--gpu_count must be -1 (all GPUs) or a positive integer, got {args.gpu_count}"
+        )
+    else:
+        gpu_count = min(args.gpu_count, available_gpus)
+    print(f"🎯 Using {gpu_count} GPU(s) out of {available_gpus} available")
+
+    if args.overlap_length >= args.context_length:
+        raise ValueError(
+            f"--overlap_length ({args.overlap_length}) must be strictly less than "
+            f"--context_length ({args.context_length})"
+        )
+
+    input_paths = resolve_input_paths(args.input)
+    if not input_paths:
+        raise ValueError("No valid input files resolved from --input.")
+    print(f"📁 Resolved {len(input_paths)} input file(s)")
+
+    os.makedirs(args.output_path, exist_ok=True)
+    run_timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    persistent_runtime: Optional[Dict[str, Any]] = None
+    interrupted = False
+    try:
+        persistent_runtime = create_persistent_worker_runtime(
+            model_name=args.model_name,
+            revision=args.revision,
+            dtype_str=dtype_str,
+            attn_implementation=args.attn_implementation,
+            gpu_count=gpu_count,
+            max_length=args.context_length,
+            overlap_length=args.overlap_length,
+            micro_batch_size=args.batch_size,
+        )
+
+        for input_idx, input_path in enumerate(input_paths, start=1):
+            print(f"\n{'=' * 80}")
+            print(f"📂 Processing input {input_idx}/{len(input_paths)}: {input_path}")
+            print(f"{'=' * 80}")
+
+            records = read_input_records(input_path, limit=args.limit)
+            if not records:
+                print("⚠️ No sequences found, skipping.")
+                continue
+
+            probs_per_head = infer_records(
+                records=records,
+                model_name=args.model_name,
+                dtype_str=dtype_str,
+                gpu_count=gpu_count,
+                max_length=args.context_length,
+                overlap_length=args.overlap_length,
+                micro_batch_size=args.batch_size,
+                dispatch_chunk_size=args.dispatch_chunk_size,
+                persistent_runtime=persistent_runtime,
+                attn_implementation=args.attn_implementation,
+                revision=args.revision,
+            )
+
+            head_names = (
+                HEAD_NAMES
+                if len(probs_per_head) == 2
+                else [f"head_{i}" for i in range(len(probs_per_head))]
+            )
+            output_parquet_path = make_output_parquet_path(
+                args.output_path, input_path, run_timestamp
+            )
+            write_prob_parquet(
+                records=records,
+                probs_per_head=probs_per_head,
+                head_names=head_names,
+                output_path=output_parquet_path,
+                prob_dtype=args.prob_dtype,
+                save_non_coding_prob=args.save_non_coding_prob,
+            )
+    except KeyboardInterrupt:
+        interrupted = True
+        raise
+    finally:
+        if persistent_runtime is not None:
+            print("🛑 Shutting down persistent GPU workers...")
+            shutdown_persistent_worker_runtime(
+                persistent_runtime, interrupted=interrupted
+            )
+
+    total_time = time.time() - total_start_time
+    minutes, seconds = divmod(total_time, 60)
+    print(f"\n⏱️ Total execution time: {int(minutes)}m {seconds:.2f}s")
+    print("✨ Completed successfully! ✨\n")
+
+
+if __name__ == "__main__":
+    mp.set_start_method("spawn", force=True)
+    main()

annotation_pipeline/src/post_processing.py

@@ -0,0 +1,511 @@
+"""Post-process Carbon-A coding probabilities into GFF3, CDS FASTA, and protein FASTA.
+
+Post-processing runs three stages on each strand of each record, in
+transcript orientation:
+
+    1. The transition network (transition_network.py) reads the coding
+       probability track and scores the four permitted transitions at every
+       base: enter gene, leave gene, enter intron, leave intron.
+    2. The Viterbi algorithm (viterbi.py) combines the coding probabilities,
+       the transition probabilities, and the gene grammar into the highest-scoring
+       labeling of the strand into intergenic, exon, and intron bases. Runs of
+       exon bases form CDS intervals, intervals separated by an intron belong
+       to one gene, and the reading frame gives each interval its phase.
+    3. The confidence network (confidence_network.py) reads the coding
+       probabilities together with the decoded roles of the whole strand,
+       neighbouring genes included, and returns, for every complete gene, the
+       probability that its structure is exactly right.
+
+The score of a labeling is the sum over bases of log P(CDS) for exon bases and
+log(1 - P(CDS)) for all other bases, plus the log-odds of the transition at
+every base where a segment begins or a gene closes, plus a fixed log-prior for
+each intron's splice class. The grammar decides which labelings are admissible:
+a gene opens at an ATG start codon and closes at the first in-frame stop codon
+of the selected genetic code, an intron opens at a donor site and closes at an
+acceptor site of a permitted class, and the reading frame continues across
+each intron. The implementation expands the three labels into 30 states so
+that every condition can be checked from local context.
+
+Output filtering:
+
+    FASTA records shorter than min_record_length are skipped before sequence
+    encoding and Viterbi decoding. The production default is 100,000 bp;
+    zero disables this record-level filter.
+
+    Gene models whose whole spliced CDS is shorter than ``min_cds_length`` are
+    omitted after decoding. The production default is
+    ``MIN_CDS_LENGTH_BP`` (96 bp); zero disables this gene-level filter.
+    This output filter does not constrain individual exon lengths or alter the
+    Viterbi path.
+
+Gene confidence:
+
+    Every complete gene carries ``confidence``, the confidence network's
+    estimate of the probability that its structure is exactly correct, in its
+    GFF3 gene attributes and FASTA headers. Genes truncated by a record edge
+    are not scored and carry ``confidence=NA``. All genes are written;
+    filtering at a threshold is left to downstream tools such as
+    evaluation.py, whose gene-level results use a threshold of 0.1.
+"""
+
+import argparse
+import os
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from constants import (
+    CHUNK_SIZE_MBP,
+    CODON_TABLE_PATH,
+    CONFIDENCE_THRESHOLD,
+    CONFIDENCE_NETWORK_PATH,
+    MAX_BOUNDARY_INTRON_LENGTH_BP,
+    MIN_CDS_LENGTH_BP,
+    MIN_RECORD_LENGTH_BP,
+    NEGATIVE_STRAND_PROBABILITY_COLUMN,
+    PARQUET_BATCH_ROWS,
+    POSITIVE_STRAND_PROBABILITY_COLUMN,
+    RECORD_ID_COLUMN,
+    SEQUENCE_COLUMN,
+    TRANSITION_NETWORK_PATH,
+    WARNING_DIRECTORY_SUFFIX,
+    WARNINGS_FILE_NAME,
+)
+from utils import (
+    load_codon_table, stop_codons_from_table,
+    probability_vector_from_parquet,
+    validate_max_boundary_intron_length,
+    encode_both_strands,
+    assemble_gene_model,
+    codon_table_mismatch,
+    write_gff3, write_cds_fasta, write_protein_fasta,
+)
+from confidence_network import load_confidence_network
+from transition_network import load_transition_network
+from viterbi import choose_viterbi_backend, decode_paths, iter_gene_exons_with_states
+
+
+def default_output_prefix(input_path):
+    """Return the default output prefix for one raw-probability parquet file."""
+    if input_path.endswith(".raw_probs.parquet"):
+        return input_path[:-len(".raw_probs.parquet")] + ".predicted"
+    return os.path.splitext(input_path)[0] + ".predicted"
+
+
+def iter_probability_chunks(
+    parquet_file, record_id_col, sequence_col,
+    positive_col, negative_col, max_chunk_bp,
+    min_record_length=0, filter_stats=None,
+):
+    """Read and encode strand tracks in approximately ``max_chunk_bp`` chunks.
+
+    Each record yields two tracks in transcript orientation, ``(coding
+    probabilities, encoded bases)`` for the forward strand and for the reverse
+    complement, together with a ``(record_id, sequence, strand, length)``
+    context. A sequence record is never split, and each yielded base count
+    includes both tracks.
+    """
+    required_columns = [record_id_col, sequence_col, positive_col, negative_col]
+    tracks = []
+    track_contexts = []
+    chunk_bp = 0
+
+    for batch in parquet_file.iter_batches(batch_size=PARQUET_BATCH_ROWS, columns=required_columns):
+        columns_by_name = {
+            column_name: batch.column(column_index)
+            for column_index, column_name in enumerate(batch.schema.names)
+        }
+        record_ids = columns_by_name[record_id_col]
+        sequences = columns_by_name[sequence_col]
+        positive_probabilities = columns_by_name[positive_col]
+        negative_probabilities = columns_by_name[negative_col]
+
+        for row_index in range(batch.num_rows):
+            sequence = sequences[row_index].as_py().upper()
+            record_length = len(sequence)
+            record_id = record_ids[row_index].as_py()
+            if record_length < min_record_length:
+                if filter_stats is not None:
+                    filter_stats["record_count"] += 1
+                    filter_stats["bp"] += record_length
+                continue
+            positive = probability_vector_from_parquet(positive_probabilities[row_index])
+            negative = probability_vector_from_parquet(negative_probabilities[row_index])
+            if len(positive) < record_length or len(negative) < record_length:
+                raise ValueError(
+                    f"record {record_id!r} has {record_length} bases but "
+                    f"probability lengths {len(positive)} and {len(negative)}"
+                )
+
+            positive = positive[:record_length]
+            negative = negative[:record_length]
+            forward_bases, reverse_bases = encode_both_strands(sequence)
+            tracks.append((positive, forward_bases))
+            track_contexts.append((record_id, sequence, "+", record_length))
+            tracks.append((np.ascontiguousarray(negative[::-1]), reverse_bases))
+            track_contexts.append((record_id, sequence, "-", record_length))
+            chunk_bp += 2 * record_length
+
+            if chunk_bp >= max_chunk_bp:
+                yield tracks, track_contexts, chunk_bp
+                tracks = []
+                track_contexts = []
+                chunk_bp = 0
+
+    if tracks:
+        yield tracks, track_contexts, chunk_bp
+
+
+def iter_prefetched_chunks(chunks):
+    """Prepare one chunk ahead while the caller decodes the current chunk.
+
+    The one-item lookahead overlaps Parquet decoding with GPU/CPU work while
+    keeping memory use bounded to one current and one prepared chunk.
+    """
+    chunk_iterator = iter(chunks)
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        next_chunk = executor.submit(next, chunk_iterator, None)
+        while True:
+            chunk = next_chunk.result()
+            if chunk is None:
+                return
+            next_chunk = executor.submit(next, chunk_iterator, None)
+            yield chunk
+
+
+def decode_chunk_genes(
+    tracks, track_contexts, transition_probabilities, viterbi_backend, stop_codons,
+    max_boundary_intron_length, min_cds_length, confidence_network,
+):
+    """Decode one chunk under one stop-codon set and score its complete genes.
+
+    Returns ``(chunk_genes, short_gene_count, viterbi_seconds,
+    confidence_seconds)``; each entry of ``chunk_genes`` is ``(gene,
+    probabilities, oriented_exons)``.
+    """
+    stage_started = time.perf_counter()
+    state_paths = decode_paths(
+        tracks,
+        transition_probabilities,
+        viterbi_backend,
+        stop_taa="TAA" in stop_codons,
+        stop_tag="TAG" in stop_codons,
+        stop_tga="TGA" in stop_codons,
+        max_boundary_intron_length=max_boundary_intron_length,
+    )
+    viterbi_seconds = time.perf_counter() - stage_started
+    if len(state_paths) != len(track_contexts):
+        raise RuntimeError(
+            f"{viterbi_backend} Viterbi backend returned {len(state_paths)} paths "
+            f"for {len(track_contexts)} tracks"
+        )
+
+    # Gene models of the chunk, each with the oriented track and exons it
+    # was decoded from, grouped by strand: the confidence network reads each
+    # strand's decoded structure whole.
+    chunk_genes = []
+    strand_gene_indices = []
+    short_gene_count = 0
+    for state_path, track_context, (probabilities, _) in zip(state_paths, track_contexts, tracks):
+        record_id, sequence, strand, record_length = track_context
+        if len(state_path) != record_length:
+            raise RuntimeError(
+                f"{viterbi_backend} Viterbi backend returned path length {len(state_path)} "
+                f"for record {record_id!r} {strand} with {record_length} bases"
+            )
+        gene_indices = []
+        for oriented_exons, first_cds_state, last_cds_state in iter_gene_exons_with_states(state_path):
+            gene = assemble_gene_model(
+                record_id, sequence, oriented_exons, strand, record_length,
+                first_cds_state, last_cds_state,
+                stop_codons,
+            )
+            if gene is None:
+                continue
+            if len(gene["cds"]) < min_cds_length:
+                short_gene_count += 1
+                continue
+            gene["confidence"] = None
+            gene_indices.append(len(chunk_genes))
+            chunk_genes.append((gene, probabilities, oriented_exons))
+        strand_gene_indices.append((probabilities, gene_indices))
+
+    # Confidence of every complete gene of the chunk, strand by strand. Every
+    # written gene gives roles to its strand; only the complete ones are scored.
+    stage_started = time.perf_counter()
+    for probabilities, gene_indices in strand_gene_indices:
+        scored = [
+            position for position, gene_index in enumerate(gene_indices)
+            if chunk_genes[gene_index][0]["complete"]
+        ]
+        if not scored:
+            continue
+        confidences = confidence_network.inference(
+            probabilities, [chunk_genes[gene_index][2] for gene_index in gene_indices], scored,
+        )
+        for position, confidence in zip(scored, confidences):
+            chunk_genes[gene_indices[position]][0]["confidence"] = float(confidence)
+    confidence_seconds = time.perf_counter() - stage_started
+    return chunk_genes, short_gene_count, viterbi_seconds, confidence_seconds
+
+
+def record_codon_table_warning(output_prefix, warning, stop_codons):
+    """Record the outcome of the codon-table check where a batch run will see it.
+
+    A failed check (``warning`` is the message) is appended to
+    ``WARNINGS_FILE_NAME`` in the parent of the output directory, one line per
+    run with the time and the directory name, and the output directory is
+    renamed with ``WARNING_DIRECTORY_SUFFIX``, so a batch run shows every
+    doubtful genome in one file and in the directory names. A passed check in
+    a directory that carries the suffix, typically a rerun with the right
+    codon table, renames it back and records that too. Returns the output
+    prefix after any rename. The directory is left alone when the new name is
+    taken or when it is the working directory.
+    """
+    output_dir = os.path.dirname(os.path.abspath(output_prefix))
+    clean_dir = output_dir[:-len(WARNING_DIRECTORY_SUFFIX)] if output_dir.endswith(WARNING_DIRECTORY_SUFFIX) else output_dir
+    flagged_dir = clean_dir + WARNING_DIRECTORY_SUFFIX
+    if warning is not None:
+        entry = warning
+        target = flagged_dir
+    elif output_dir == flagged_dir:
+        entry = f"cleared: the codon-table check passed with stops={'+'.join(stop_codons)}"
+        target = clean_dir
+    else:
+        return output_prefix
+
+    with open(os.path.join(os.path.dirname(clean_dir), WARNINGS_FILE_NAME), "a") as handle:
+        handle.write(f"{time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}\t{os.path.basename(clean_dir)}\t{entry}\n")
+    if target == output_dir:
+        return output_prefix
+    if os.path.exists(target) or output_dir == os.getcwd():
+        print(
+            f"[post_processing] leaving {output_dir} in place: {target} exists or it is the working directory",
+            file=sys.stderr, flush=True,
+        )
+        return output_prefix
+    os.rename(output_dir, target)
+    return os.path.join(target, os.path.basename(output_prefix))
+
+
+def post_process_probability_file(
+    input_path, output_prefix, viterbi_backend, codon_to_amino_acid,
+    transition_network, confidence_network,
+    max_boundary_intron_length, chunk_mbp,
+    min_record_length=MIN_RECORD_LENGTH_BP,
+    min_cds_length=MIN_CDS_LENGTH_BP,
+):
+    """Post-process one raw-probability parquet into GFF3, CDS FASTA, and protein FASTA.
+
+    ``viterbi_backend`` is the name returned by ``choose_viterbi_backend``.
+    When the confidences of the complete genes contradict the codon table
+    (utils.codon_table_mismatch), the warning is printed and recorded in the
+    directory layout by ``record_codon_table_warning``.
+    """
+    import pyarrow.parquet as pq
+
+    file_started = time.perf_counter()
+    stop_codons = tuple(stop_codons_from_table(codon_to_amino_acid))
+    parquet_file = pq.ParquetFile(input_path)
+    required_columns = {
+        RECORD_ID_COLUMN,
+        SEQUENCE_COLUMN,
+        POSITIVE_STRAND_PROBABILITY_COLUMN,
+        NEGATIVE_STRAND_PROBABILITY_COLUMN,
+    }
+    missing_columns = required_columns - set(parquet_file.schema_arrow.names)
+    if missing_columns:
+        raise ValueError(
+            "raw probability parquet is missing columns: "
+            + ", ".join(sorted(missing_columns))
+        )
+
+    genes = []
+    short_gene_count = 0
+    total_bp = 0
+    transition_seconds = 0.0
+    viterbi_seconds = 0.0
+    confidence_seconds = 0.0
+    filter_stats = {"record_count": 0, "bp": 0}
+
+    input_chunks = iter_probability_chunks(
+        parquet_file, RECORD_ID_COLUMN, SEQUENCE_COLUMN,
+        POSITIVE_STRAND_PROBABILITY_COLUMN, NEGATIVE_STRAND_PROBABILITY_COLUMN,
+        chunk_mbp * 1_000_000,
+        min_record_length=min_record_length,
+        filter_stats=filter_stats,
+    )
+    for tracks, track_contexts, chunk_bp in iter_prefetched_chunks(input_chunks):
+        # Stage 1: transition probabilities for every oriented track of the chunk.
+        stage_started = time.perf_counter()
+        transition_probabilities = [
+            transition_network.inference(probabilities)
+            for probabilities, _ in tracks
+        ]
+        transition_seconds += time.perf_counter() - stage_started
+
+        # Stages 2 and 3: the Viterbi algorithm, gene models, and their confidences.
+        chunk_genes, chunk_short_gene_count, chunk_viterbi_seconds, chunk_confidence_seconds = decode_chunk_genes(
+            tracks, track_contexts, transition_probabilities, viterbi_backend, stop_codons,
+            max_boundary_intron_length, min_cds_length, confidence_network,
+        )
+        short_gene_count += chunk_short_gene_count
+        viterbi_seconds += chunk_viterbi_seconds
+        confidence_seconds += chunk_confidence_seconds
+
+        genes.extend(gene for gene, _, _ in chunk_genes)
+        total_bp += chunk_bp
+
+    genes.sort(key=lambda gene: (gene["record_id"], gene["exons"][0][0]))
+    write_gff3(genes, output_prefix + ".gff")
+    write_cds_fasta(genes, output_prefix + ".fna")
+    protein_count = write_protein_fasta(genes, output_prefix + ".faa", codon_to_amino_acid)
+
+    total_seconds = time.perf_counter() - file_started
+    rate = total_bp / viterbi_seconds / 1e6 if viterbi_seconds else 0.0
+    confidences = sorted(gene["confidence"] for gene in genes if gene["confidence"] is not None)
+    scored_gene_count = len(confidences)
+    median_confidence = confidences[len(confidences) // 2] if confidences else float("nan")
+    confident_gene_count = sum(confidence >= CONFIDENCE_THRESHOLD for confidence in confidences)
+    warning = None
+    if codon_table_mismatch(confidences):
+        warning = (
+            f"WARNING: only {confident_gene_count} of {scored_gene_count} complete genes "
+            f"({confident_gene_count / scored_gene_count:.0%}) reach confidence {CONFIDENCE_THRESHOLD} "
+            f"with stops={'+'.join(stop_codons)}. "
+            f"The codon table probably does not match this genome; check --codon_table."
+        )
+        print(f"[post_processing] {warning}", file=sys.stderr, flush=True)
+    output_prefix = record_codon_table_warning(output_prefix, warning, stop_codons)
+    print(
+        f"[post_processing] viterbi_backend={viterbi_backend}  stops={'+'.join(stop_codons)}  "
+        f"max_boundary_intron_length={max_boundary_intron_length}  "
+        f"min_record_length={min_record_length}  "
+        f"{os.path.basename(input_path)}: {total_bp / 1e6:.0f} Mbp; "
+        f"transition network {transition_seconds:.1f}s; "
+        f"Viterbi {viterbi_seconds:.1f}s ({rate:.0f} Mbp/s); "
+        f"confidence network {confidence_seconds:.1f}s "
+        f"({scored_gene_count} genes scored, {confident_gene_count} at or above {CONFIDENCE_THRESHOLD}, "
+        f"median confidence {median_confidence:.3f}); "
+        f"{total_seconds:.1f}s end to end  ->  {len(genes)} genes  "
+        f"({short_gene_count} shorter than {min_cds_length} bp omitted)  "
+        f"({filter_stats['record_count']} records / "
+        f"{filter_stats['bp'] / 1e6:.1f} Mbp shorter than "
+        f"{min_record_length} bp skipped)  "
+        f"->  {output_prefix}.gff/.fna/.faa ({protein_count} proteins)",
+        flush=True,
+    )
+
+
+def parse_args():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--input", required=True, help="one *.raw_probs.parquet file")
+    parser.add_argument("--output", default=None, help="output prefix; defaults to <input>.predicted")
+    parser.add_argument(
+        "--viterbi_backend",
+        choices=["auto", "gpu", "cpu", "numba"],
+        default="auto",
+        help="Viterbi backend: gpu (CUDA), cpu (C++/OpenMP), or numba (reference); auto prefers gpu",
+    )
+    parser.add_argument(
+        "--codon_table",
+        default=str(CODON_TABLE_PATH),
+        help=(
+            "JSON object mapping all 64 ATCG codons to amino acids; '*' marks "
+            "decoder stop codons"
+        ),
+    )
+    parser.add_argument(
+        "--transition_network",
+        default=str(TRANSITION_NETWORK_PATH),
+        help="transition network weights (default: weights/transition_network.safetensors)",
+    )
+    parser.add_argument(
+        "--confidence_network",
+        default=str(CONFIDENCE_NETWORK_PATH),
+        help="confidence network weights (default: weights/confidence_network.safetensors)",
+    )
+    parser.add_argument(
+        "--device",
+        default=None,
+        help="torch device for the two networks (default: cuda if available, else cpu)",
+    )
+    parser.add_argument(
+        "--max_boundary_intron_length",
+        type=int,
+        default=MAX_BOUNDARY_INTRON_LENGTH_BP,
+        help=(
+            "maximum visible boundary-intron length in bp; -1 is unlimited; "
+            "0 disables boundary introns, leaving only exon-truncated incomplete genes"
+        ),
+    )
+    parser.add_argument(
+        "--chunk_mbp",
+        type=int,
+        default=CHUNK_SIZE_MBP,
+        help="target host-side chunk size in megabases; sequence records are never split",
+    )
+    parser.add_argument(
+        "--min_record_length",
+        type=int,
+        default=MIN_RECORD_LENGTH_BP,
+        help=(
+            "minimum input FASTA record length in bp; shorter records are skipped "
+            "before decoding (0 disables)"
+        ),
+    )
+    parser.add_argument(
+        "--min_cds_length",
+        type=int,
+        default=MIN_CDS_LENGTH_BP,
+        help=(
+            "minimum spliced CDS length in bp for output gene models; "
+            "shorter models are omitted (0 disables)"
+        ),
+    )
+    args = parser.parse_args()
+    try:
+        args.max_boundary_intron_length = validate_max_boundary_intron_length(
+            args.max_boundary_intron_length
+        )
+    except (TypeError, ValueError) as exc:
+        parser.error(str(exc))
+    if args.min_record_length < 0:
+        parser.error("--min_record_length must be non-negative")
+    if args.min_cds_length < 0:
+        parser.error("--min_cds_length must be non-negative")
+    return args
+
+
+def main():
+    """Command-line entry point."""
+    args = parse_args()
+    try:
+        codon_to_amino_acid = load_codon_table(args.codon_table)
+    except OSError as exc:
+        raise SystemExit(f"error: cannot read codon table {args.codon_table}: {exc}") from exc
+    except ValueError as exc:
+        raise SystemExit(f"error: {exc}") from exc
+
+    output_prefix = args.output or default_output_prefix(args.input)
+    viterbi_backend = choose_viterbi_backend(args.viterbi_backend)
+    device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    transition_network = load_transition_network(args.transition_network, device)
+    confidence_network = load_confidence_network(args.confidence_network, device)
+    post_process_probability_file(
+        args.input, output_prefix, viterbi_backend, codon_to_amino_acid,
+        transition_network, confidence_network,
+        max_boundary_intron_length=args.max_boundary_intron_length,
+        chunk_mbp=args.chunk_mbp,
+        min_record_length=args.min_record_length,
+        min_cds_length=args.min_cds_length,
+    )
+
+
+if __name__ == "__main__":
+    main()

annotation_pipeline/src/tcn.py

@@ -0,0 +1,136 @@
+"""Temporal convolutional network shared by the transition and confidence networks.
+
+Both networks of the post-processing pipeline read per-base tracks and need
+one contextual feature vector per base. They use the same architecture: a
+width-5 input convolution, a stack of residual blocks with depthwise dilated
+convolutions (dilations 1, 2, 4, ...), and a final LayerNorm. With eight
+blocks the receptive radius is 2 + 2 * (1 + 2 + ... + 128) = 512 bases.
+
+Both networks take the coding-probability track itself as input and turn it
+into log-odds internally (``coding_log_odds``). Both read tracks far longer
+than GPU memory allows in one pass; ``chunked_inference`` runs a network
+over a track in chunks with the network's full context on both sides, so the
+outputs equal those of one pass over the whole track. Both store their weights
+as safetensors files, read by ``load_weights``.
+"""
+
+import json
+
+import numpy as np
+import torch
+from safetensors import safe_open
+from torch import nn
+
+from constants import NETWORK_CHUNK_BP, NETWORK_INPUT_MIN_PROB
+
+
+class ResidualDilatedBlock(nn.Module):
+    """Add one depthwise dilated-convolution update to the features.
+
+    Each block widens the receptive field by 2 * dilation bases on both sides:
+    a depthwise width-5 dilated convolution mixes nearby positions, LayerNorm
+    and GELU shape the update, and a pointwise convolution mixes channels
+    before the residual sum.
+    """
+
+    def __init__(self, channels, dilation):
+        super().__init__()
+        self.depthwise = nn.Conv1d(
+            channels, channels, kernel_size=5,
+            padding=2 * dilation, dilation=dilation, groups=channels,
+        )
+        self.normalization = nn.LayerNorm(channels)
+        self.pointwise = nn.Conv1d(channels, channels, kernel_size=1)
+        self.activation = nn.GELU()
+
+    def forward(self, features):
+        update = self.depthwise(features)
+        update = self.normalization(update.transpose(1, 2)).transpose(1, 2)
+        update = self.activation(update)
+        return features + self.pointwise(update)
+
+
+class TCN(nn.Module):
+    """Contextual per-base features from a few input tracks."""
+
+    def __init__(self, input_channels, channels, dilation_count):
+        super().__init__()
+        self.channels = int(channels)
+        self.dilation_count = int(dilation_count)
+        if self.dilation_count < 1:
+            raise ValueError("at least one dilated block is required")
+        # Block k dilates by 2 ** k, so the receptive radius doubles with every block.
+        self.dilations = tuple(2 ** index for index in range(self.dilation_count))
+        self.context_radius = 2 + 2 * sum(self.dilations)
+
+        self.input_projection = nn.Conv1d(input_channels, self.channels, kernel_size=5, padding=2)
+        self.blocks = nn.Sequential(*[
+            ResidualDilatedBlock(self.channels, dilation) for dilation in self.dilations
+        ])
+        self.output_normalization = nn.LayerNorm(self.channels)
+
+    def forward(self, input_tracks):
+        """Map ``[batch, input_channels, bases]`` to ``[batch, bases, channels]``."""
+        features = self.blocks(self.input_projection(input_tracks))
+        return self.output_normalization(features.transpose(1, 2))
+
+
+def coding_log_odds(probabilities):
+    """Return the float32 log-odds of a tensor of coding probabilities.
+
+    Probabilities are clipped to ``[NETWORK_INPUT_MIN_PROB, 1 - NETWORK_INPUT_MIN_PROB]``
+    first, so a probability of exactly 0 or 1 stays finite.
+    """
+    return torch.logit(probabilities.float().clamp(NETWORK_INPUT_MIN_PROB, 1.0 - NETWORK_INPUT_MIN_PROB))
+
+
+def chunked_inference(network, context_radius, tracks, device, chunk_bp=NETWORK_CHUNK_BP):
+    """Run a per-base network over long tracks, one chunk at a time.
+
+    ``tracks`` are per-base arrays of one common length in decoder
+    orientation. ``network(*chunks)`` receives corresponding chunks of them as
+    tensors on ``device`` and returns one output row per base; each output
+    depends on the inputs within ``context_radius`` bases of its base. The
+    tracks are processed in chunks of at most ``chunk_bp`` bases, each
+    evaluated with ``context_radius`` extra bases on both sides, clipped to
+    the tracks, so the outputs equal those of one pass over the whole tracks
+    while only one chunk is ever on the device. Returns the float32
+    ``[bases, ...]`` outputs as one CPU array. On CUDA the network runs under
+    bfloat16 autocast, the precision used for all reported results; on CPU it
+    runs in float32.
+    """
+    tracks = [np.asarray(track) for track in tracks]
+    if not tracks or any(track.ndim != 1 for track in tracks) or len({len(track) for track in tracks}) != 1:
+        raise ValueError("tracks must be one or more vectors of one common length")
+    length = len(tracks[0])
+    if length == 0:
+        raise ValueError("tracks must be nonempty")
+    if chunk_bp < 1:
+        raise ValueError("chunk size must be positive")
+
+    outputs = None
+    for start in range(0, length, chunk_bp):
+        stop = min(start + chunk_bp, length)
+        window_start = max(0, start - context_radius)
+        window_stop = min(length, stop + context_radius)
+        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
+            window_outputs = network(*(
+                torch.as_tensor(track[window_start:window_stop]).to(device) for track in tracks
+            ))
+        chunk_outputs = window_outputs[start - window_start:stop - window_start].float().cpu().numpy()
+        if outputs is None:
+            outputs = np.empty((length,) + chunk_outputs.shape[1:], dtype=np.float32)
+        outputs[start:stop] = chunk_outputs
+    return outputs
+
+
+def load_weights(path):
+    """Read a network's weights file: its metadata fields plus ``state_dict``.
+
+    The weights are a safetensors file whose metadata values are JSON-encoded,
+    such as the network type, its sizes, and the names of its inputs or outputs.
+    """
+    with safe_open(path, framework="pt") as weights:
+        checkpoint = {key: json.loads(value) for key, value in (weights.metadata() or {}).items()}
+        checkpoint["state_dict"] = {name: weights.get_tensor(name) for name in weights.keys()}
+    return checkpoint

annotation_pipeline/src/transition_network.py

@@ -0,0 +1,141 @@
+"""Transition network: transition probabilities between adjacent bases.
+
+The transition network is a small temporal convolutional network (TCN) that
+reads one strand's coding-probability track and predicts what happens at
+every boundary between adjacent bases and at both ends of the record. An
+L-base record has L + 1 such boundaries; boundary j lies between bases j - 1
+and j. At each boundary the network outputs a distribution over five
+possibilities, staying or one of the four transitions:
+
+    stay            remain in the current segment
+    enter_gene      a gene begins with base j
+    leave_gene      a gene ends with base j - 1
+    enter_intron    an intron begins with base j
+    leave_intron    an intron ends and an exon resumes with base j
+
+The Viterbi algorithm (viterbi.py) adds the log-probability of the
+possibility it selects at every boundary to the path score. Because every path
+selects exactly one per boundary, subtracting log P(stay) at each
+boundary changes the score of every path by the same constant, so the decoders
+receive each transition as its log-odds relative to staying, and staying costs
+nothing. The four log-odds are attached to the base where the new segment
+begins or, for leaving a gene, to the base that closes it:
+
+    enter_gene(i)   = log P(enter_gene   | boundary i)     - log P(stay | boundary i)
+    leave_gene(i)   = log P(leave_gene   | boundary i + 1) - log P(stay | boundary i + 1)
+    enter_intron(i) = log P(enter_intron | boundary i)     - log P(stay | boundary i)
+    leave_intron(i) = log P(leave_intron | boundary i)     - log P(stay | boundary i)
+
+Records are scored in chunks read with the network's full context on both
+sides (tcn.chunked_inference), so a chromosome never has to fit in GPU memory
+at once. The scores leave this module in nats; the Viterbi implementations
+quantize them to their own fixed-point units (utils.prepare_transition_probabilities).
+"""
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from constants import (
+    ENTER_GENE,
+    ENTER_INTRON,
+    LEAVE_GENE,
+    LEAVE_INTRON,
+    NETWORK_CHUNK_BP,
+    TRANSITION_COUNT,
+    TRANSITION_NAMES,
+)
+from tcn import TCN, chunked_inference, coding_log_odds, load_weights
+
+TRANSITION_NETWORK_TYPE = "transition_network"
+
+# Columns of the network's boundary logits: staying, then the four transitions
+# in constants.TRANSITION_NAMES order.
+STAY = 0
+TRANSITION_LOGIT_NAMES = ("stay",) + TRANSITION_NAMES
+
+
+class TransitionNetwork(nn.Module):
+    """Predict the five-way transition distribution at every boundary.
+
+    The TCN reads the coding probabilities, as log-odds, alone. A width-two
+    readout then combines the features of the two bases beside each boundary,
+    so L input bases yield L + 1 distributions; the boundaries before
+    the first and after the last base see zero features on their outer side.
+    The context radius is one base larger than the TCN's because a boundary
+    reads both neighbours.
+    """
+
+    def __init__(self, channels=32, dilation_count=8):
+        super().__init__()
+        self.encoder = TCN(1, channels, dilation_count)
+        self.context_radius = self.encoder.context_radius + 1
+        self.transition_readout = nn.Conv1d(self.encoder.channels, len(TRANSITION_LOGIT_NAMES), kernel_size=2)
+
+    def transition_logits(self, probabilities):
+        """Return ``[batch, bases + 1, 5]`` boundary logits over ``TRANSITION_LOGIT_NAMES`` from ``[batch, bases]`` probabilities."""
+        if probabilities.ndim != 2:
+            raise ValueError("probabilities must have shape [batch, bases]")
+        features = self.encoder(coding_log_odds(probabilities)[:, None, :]).transpose(1, 2)
+        return self.transition_readout(F.pad(features, (1, 1))).transpose(1, 2)
+
+    def forward(self, probabilities):
+        """Return the four per-base transition tracks, in nats, from ``[batch, bases]`` probabilities."""
+        return transition_log_odds_from_logits(self.transition_logits(probabilities))
+
+    @torch.inference_mode()
+    def inference(self, probabilities, chunk_bp=NETWORK_CHUNK_BP):
+        """Return the transition probabilities of one oriented coding-probability track.
+
+        The result is a float32 ``(bases, 4)`` array in ``TRANSITION_NAMES``
+        column order, each value the transition's log-odds against staying, in
+        nats. The track is processed in ``chunk_bp`` chunks with the network's
+        full context on both sides (tcn.chunked_inference), so the values equal
+        those of scoring the whole record at once.
+        """
+        probabilities = np.asarray(probabilities, dtype=np.float32)
+        if probabilities.ndim != 1 or len(probabilities) == 0:
+            raise ValueError("probabilities must be a nonempty vector")
+
+        def network(chunk):
+            return self(chunk[None, :])[0]
+
+        return chunked_inference(
+            network, self.context_radius, (probabilities,), next(self.parameters()).device, chunk_bp,
+        )
+
+
+def transition_log_odds_from_logits(transition_logits):
+    """Map ``[batch, bases + 1, 5]`` boundary logits to ``[batch, bases, 4]`` tracks.
+
+    Each track is the transition's log-odds relative to ``stay`` at the same
+    boundary, placed on the base described in the module docstring: the
+    boundary before the base for entering a gene or an intron and for leaving
+    an intron, the boundary after the base for leaving a gene.
+    """
+    if transition_logits.ndim != 3 or transition_logits.shape[2] != len(TRANSITION_LOGIT_NAMES):
+        raise ValueError("transition logits must have shape [batch, bases + 1, 5]")
+    if transition_logits.shape[1] < 2:
+        raise ValueError("transition logits must describe at least one base")
+
+    # Log-odds against staying; the transition columns follow the stay column.
+    relative = (transition_logits - transition_logits[:, :, STAY:STAY + 1])[:, :, STAY + 1:]
+    tracks = relative.new_zeros(relative.shape[0], relative.shape[1] - 1, TRANSITION_COUNT)
+    tracks[:, :, ENTER_GENE] = relative[:, :-1, ENTER_GENE]
+    tracks[:, :, LEAVE_GENE] = relative[:, 1:, LEAVE_GENE]
+    tracks[:, :, ENTER_INTRON] = relative[:, :-1, ENTER_INTRON]
+    tracks[:, :, LEAVE_INTRON] = relative[:, :-1, LEAVE_INTRON]
+    return tracks
+
+
+def load_transition_network(path, device):
+    """Load the transition network weights onto ``device`` in evaluation mode."""
+    checkpoint = load_weights(path)
+    if checkpoint.get("model_type") != TRANSITION_NETWORK_TYPE:
+        raise ValueError(f"{path}: not a transition network checkpoint")
+    if tuple(checkpoint.get("outcome_names", TRANSITION_LOGIT_NAMES)) != TRANSITION_LOGIT_NAMES:
+        raise ValueError(f"{path}: unexpected transition logits")
+    network = TransitionNetwork(int(checkpoint["channels"]), int(checkpoint["dilation_count"]))
+    network.load_state_dict(checkpoint["state_dict"])
+    return network.to(device).eval()

annotation_pipeline/src/utils.py

@@ -0,0 +1,791 @@
+"""Utility functions shared by the Carbon-A annotation pipeline.
+
+The functions in this module do not import numba, torch, or Viterbi backends.
+They are intentionally plain NumPy/Python helpers used by post-processing,
+evaluation, and the CPU/GPU Viterbi backends.
+"""
+
+import json
+import operator
+from pathlib import Path
+
+import numpy as np
+
+from constants import (
+    AMINO_ACID_SYMBOLS,
+    BASE_A,
+    BASE_C,
+    BASE_G,
+    BASE_N,
+    BASE_T,
+    BASE_VALUES,
+    CODON_BASES,
+    CODON_TABLE_WARNING_MIN_GENES,
+    CODON_TABLE_WARNING_FRACTION,
+    CONFIDENCE_THRESHOLD,
+    COMPLEMENT_BASES,
+    FIXED_POINT_SCORE_SCALE,
+    MIN_INTERNAL_INTRON_LENGTH_BP,
+    NEGATIVE_STRAND_PROBABILITY_COLUMN,
+    NON_NUCLEAR_ASSEMBLY_UNIT,
+    POSITIVE_STRAND_PROBABILITY_COLUMN,
+    RECORD_ID_COLUMN,
+    SEQUENCE_COLUMN,
+    START_CODONS,
+    STOP_CODONS,
+    TRANSITION_COUNT,
+)
+
+_BASE_ENCODING_VALUES = bytearray([BASE_N] * 256)
+_COMPLEMENT_ENCODING_VALUES = bytearray([BASE_N] * 256)
+for _base, _base_value in BASE_VALUES.items():
+    _BASE_ENCODING_VALUES[ord(_base)] = _base_value
+    _BASE_ENCODING_VALUES[ord(_base.lower())] = _base_value
+
+    _complement_value = BASE_VALUES[COMPLEMENT_BASES[_base]]
+    _COMPLEMENT_ENCODING_VALUES[ord(_base)] = _complement_value
+    _COMPLEMENT_ENCODING_VALUES[ord(_base.lower())] = _complement_value
+
+BASE_ENCODING_TABLE = bytes(_BASE_ENCODING_VALUES)
+COMPLEMENT_ENCODING_TABLE = bytes(_COMPLEMENT_ENCODING_VALUES)
+REVERSE_COMPLEMENT_TRANSLATION = str.maketrans("ATCGNatcgn", "TAGCNtagcn")
+
+def validate_max_boundary_intron_length(value):
+    """Return a boundary-intron limit of -1 or a nonnegative integer.
+
+    ``-1`` leaves boundary introns unlimited, ``0`` disables both boundary
+    intron states, and a positive value sets their maximum visible length in
+    bases. With a limit of ``0``, incomplete genes can only result from a record
+    boundary truncating an exon.
+    """
+    try:
+        value = operator.index(value)
+    except TypeError as exc:
+        raise TypeError("max_boundary_intron_length must be an integer") from exc
+    if value < -1:
+        raise ValueError(
+            "max_boundary_intron_length must be -1 or a nonnegative integer"
+        )
+    return value
+
+
+def validate_reconstruction_batch_size(value):
+    """Return a reconstruction batch size in bases, or ``None`` for auto sizing."""
+    if value is None:
+        return None
+    try:
+        value = operator.index(value)
+    except TypeError as exc:
+        raise TypeError("reconstruction_batch_size must be an integer or None") from exc
+    if value <= 0:
+        raise ValueError("reconstruction_batch_size must be positive or None")
+    return value
+
+
+# Codon tables are flat JSON objects: all 64 concrete codons map to one-letter
+# amino acid symbols, and "*" selects the decoder's enabled stop codons.
+ALL_CODONS = frozenset(
+    first + second + third
+    for first in CODON_BASES
+    for second in CODON_BASES
+    for third in CODON_BASES
+)
+def validate_codon_table(codon_to_amino_acid):
+    """Validate and return a copy of a complete codon-to-amino-acid mapping."""
+    if not isinstance(codon_to_amino_acid, dict):
+        raise ValueError("codon table JSON must contain one object")
+
+    supplied_codons = set(codon_to_amino_acid)
+    missing_codons = sorted(ALL_CODONS - supplied_codons)
+    unexpected_codons = sorted(supplied_codons - ALL_CODONS)
+    if missing_codons or unexpected_codons:
+        details = []
+        if missing_codons:
+            details.append("missing " + ",".join(missing_codons))
+        if unexpected_codons:
+            details.append("unexpected " + ",".join(unexpected_codons))
+        raise ValueError(
+            "codon table must contain all 64 ATCG codons: "
+            + "; ".join(details)
+        )
+
+    invalid_amino_acids = {
+        codon: amino_acid
+        for codon, amino_acid in codon_to_amino_acid.items()
+        if not isinstance(amino_acid, str)
+        or len(amino_acid) != 1
+        or amino_acid not in AMINO_ACID_SYMBOLS
+    }
+    if invalid_amino_acids:
+        formatted_values = ", ".join(
+            f"{codon}={amino_acid!r}"
+            for codon, amino_acid in sorted(invalid_amino_acids.items())
+        )
+        raise ValueError(
+            "amino acids must use the 20 standard one-letter symbols or '*': "
+            + formatted_values
+        )
+
+    unsupported_stop_codons = sorted(
+        codon
+        for codon, amino_acid in codon_to_amino_acid.items()
+        if amino_acid == "*" and codon not in STOP_CODONS
+    )
+    if unsupported_stop_codons:
+        raise ValueError(
+            "decoder stop codons must be a subset of TAA,TAG,TGA: "
+            + ",".join(unsupported_stop_codons)
+        )
+
+    return dict(codon_to_amino_acid)
+
+
+def load_codon_table(path):
+    """Load a codon-to-amino-acid mapping from a JSON file."""
+    path = Path(path)
+    try:
+        with path.open(encoding="utf-8") as handle:
+            codon_to_amino_acid = json.load(handle)
+    except json.JSONDecodeError as exc:
+        raise ValueError(
+            f"invalid codon table JSON at {path}:{exc.lineno}:{exc.colno}: {exc.msg}"
+        ) from exc
+    return validate_codon_table(codon_to_amino_acid)
+
+
+def stop_codons_from_table(codon_to_amino_acid):
+    """Return supported stop codons whose amino-acid symbol is ``*``."""
+    codon_to_amino_acid = validate_codon_table(codon_to_amino_acid)
+    return tuple(
+        codon
+        for codon in STOP_CODONS
+        if codon_to_amino_acid[codon] == "*"
+    )
+
+
+# Translation uses a five-value alphabet: A/T/C/G plus one value shared by all
+# ambiguous bases. Three values form a base-5 codon index into this table.
+TRANSLATION_BASE_COUNT = 5
+AMBIGUOUS_TRANSLATION_BASE = BASE_N
+TRANSLATION_BASE_TO_VALUE = np.full(
+    256,
+    AMBIGUOUS_TRANSLATION_BASE,
+    dtype=np.uint8,
+)
+for _base, _base_value in BASE_VALUES.items():
+    TRANSLATION_BASE_TO_VALUE[ord(_base)] = _base_value
+    TRANSLATION_BASE_TO_VALUE[ord(_base.lower())] = _base_value
+
+
+def build_translation_lookup(codon_to_amino_acid):
+    """Build the compact lookup array used by vectorized protein translation."""
+    codon_to_amino_acid = validate_codon_table(codon_to_amino_acid)
+    lookup = np.full(
+        TRANSLATION_BASE_COUNT ** 3,
+        ord("X"),
+        dtype=np.uint8,
+    )
+    for codon, amino_acid in codon_to_amino_acid.items():
+        first_base, second_base, third_base = (
+            BASE_VALUES[base] for base in codon
+        )
+        codon_index = (
+            first_base * TRANSLATION_BASE_COUNT * TRANSLATION_BASE_COUNT
+            + second_base * TRANSLATION_BASE_COUNT
+            + third_base
+        )
+        lookup[codon_index] = ord(amino_acid)
+    return lookup
+
+
+# Decoder state schema. The active sequence-conditioned decoder uses 30 states.
+PREFIX_TYPE_COUNT = 6
+STATE_COUNT = 30
+
+# Prefix type ids carried by exon and intron states. Codon phase is the number
+# of CDS bases already consumed in the current codon. Prefix types retain the
+# additional bases needed to recognize a standard stop on the next emission.
+PREFIX_TYPE_0 = 0  # codon phase 0; codon boundary, no open prefix
+PREFIX_TYPE_1 = 1  # codon phase 1; one-base prefix T
+PREFIX_TYPE_2 = 2  # codon phase 1; one-base prefix other
+PREFIX_TYPE_3 = 3  # codon phase 2; two-base prefix TA
+PREFIX_TYPE_4 = 4  # codon phase 2; two-base prefix TG
+PREFIX_TYPE_5 = 5  # codon phase 2; two-base prefix other
+
+# Numeric state ids. Keep the layout grouped by biological role rather than by
+# historical patches; lookup tables below make prefix-type mapping explicit.
+INTERGENIC = 0
+START_A = 1
+START_AT = 2
+EXON_PREFIX_TYPE_0 = 3
+EXON_PREFIX_TYPE_1 = 4
+EXON_PREFIX_TYPE_2 = 5
+EXON_PREFIX_TYPE_3 = 6
+EXON_PREFIX_TYPE_4 = 7
+EXON_PREFIX_TYPE_5 = 8
+STOP = 9
+GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0 = 10
+GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_1 = 11
+GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_2 = 12
+GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_3 = 13
+GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_4 = 14
+GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_5 = 15
+GT_GC_AG_INTRON_READY_PREFIX_TYPE_0 = 16
+GT_GC_AG_INTRON_READY_PREFIX_TYPE_1 = 17
+GT_GC_AG_INTRON_READY_PREFIX_TYPE_2 = 18
+GT_GC_AG_INTRON_READY_PREFIX_TYPE_3 = 19
+GT_GC_AG_INTRON_READY_PREFIX_TYPE_4 = 20
+GT_GC_AG_INTRON_READY_PREFIX_TYPE_5 = 21
+AT_AC_INTRON_PREFIX_TYPE_0 = 22
+AT_AC_INTRON_PREFIX_TYPE_1 = 23
+AT_AC_INTRON_PREFIX_TYPE_2 = 24
+AT_AC_INTRON_PREFIX_TYPE_3 = 25
+AT_AC_INTRON_PREFIX_TYPE_4 = 26
+AT_AC_INTRON_PREFIX_TYPE_5 = 27
+LEADING_BOUNDARY_INTRON = 28
+TRAILING_BOUNDARY_INTRON = 29
+
+EXON_STATE_BY_PREFIX_TYPE = np.array((
+    EXON_PREFIX_TYPE_0,
+    EXON_PREFIX_TYPE_1,
+    EXON_PREFIX_TYPE_2,
+    EXON_PREFIX_TYPE_3,
+    EXON_PREFIX_TYPE_4,
+    EXON_PREFIX_TYPE_5,
+), dtype=np.int64)
+
+GT_GC_AG_INTRON_FIRST_STATE_BY_PREFIX_TYPE = np.array((
+    GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0,
+    GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_1,
+    GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_2,
+    GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_3,
+    GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_4,
+    GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_5,
+), dtype=np.int64)
+
+GT_GC_AG_INTRON_READY_STATE_BY_PREFIX_TYPE = np.array((
+    GT_GC_AG_INTRON_READY_PREFIX_TYPE_0,
+    GT_GC_AG_INTRON_READY_PREFIX_TYPE_1,
+    GT_GC_AG_INTRON_READY_PREFIX_TYPE_2,
+    GT_GC_AG_INTRON_READY_PREFIX_TYPE_3,
+    GT_GC_AG_INTRON_READY_PREFIX_TYPE_4,
+    GT_GC_AG_INTRON_READY_PREFIX_TYPE_5,
+), dtype=np.int64)
+
+AT_AC_INTRON_STATE_BY_PREFIX_TYPE = np.array((
+    AT_AC_INTRON_PREFIX_TYPE_0,
+    AT_AC_INTRON_PREFIX_TYPE_1,
+    AT_AC_INTRON_PREFIX_TYPE_2,
+    AT_AC_INTRON_PREFIX_TYPE_3,
+    AT_AC_INTRON_PREFIX_TYPE_4,
+    AT_AC_INTRON_PREFIX_TYPE_5,
+), dtype=np.int64)
+
+
+# Unlike codon phase above, GFF3 phase is the number of visible bases to skip at
+# the beginning of a CDS feature before reaching the first complete codon. A
+# boundary-truncated gene's first state retains enough information to recover it.
+INITIAL_CDS_PHASE_BY_STATE = {
+    START_A: 0,
+    START_AT: 2,
+    EXON_PREFIX_TYPE_0: 1,
+    EXON_PREFIX_TYPE_1: 0,
+    EXON_PREFIX_TYPE_2: 0,
+    EXON_PREFIX_TYPE_3: 2,
+    EXON_PREFIX_TYPE_4: 2,
+    EXON_PREFIX_TYPE_5: 2,
+    STOP: 1,
+}
+
+
+def initial_cds_phase_from_state(state):
+    """Return the initial GFF3 phase encoded by the first visible CDS state."""
+    try:
+        return INITIAL_CDS_PHASE_BY_STATE[int(state)]
+    except KeyError as exc:
+        raise ValueError(f"state {state} does not emit CDS") from exc
+
+
+def reverse_complement(sequence):
+    """Return the reverse complement string for a DNA sequence."""
+    return sequence.translate(REVERSE_COMPLEMENT_TRANSLATION)[::-1]
+
+
+def encode_both_strands(sequence):
+    """Encode both strands as A/T/C/G/N = 0/1/2/3/4."""
+    sequence_bytes = sequence.encode("ascii")
+    forward_bases = np.frombuffer(
+        sequence_bytes.translate(BASE_ENCODING_TABLE),
+        dtype=np.int8,
+    )
+    reverse_bases = np.frombuffer(
+        sequence_bytes.translate(COMPLEMENT_ENCODING_TABLE)[::-1],
+        dtype=np.int8,
+    )
+    return forward_bases, reverse_bases
+
+
+def probability_vector_from_parquet(value):
+    """Convert one PyArrow list cell to a 1D float16 probability array."""
+    if not hasattr(value, "values"):
+        raise TypeError("probability column must contain PyArrow list values")
+    probabilities = np.asarray(value.values, dtype=np.float16)
+
+    if probabilities.ndim != 1:
+        raise ValueError(f"probability vector must be 1D, got {probabilities.shape}")
+    return probabilities
+
+
+def prepare_decode_tasks(tasks):
+    """Normalize decoder tasks and reject invalid shapes or lengths."""
+    prepared_tasks = []
+    for task_index, (probabilities, encoded_bases) in enumerate(tasks):
+        probability_array = np.ascontiguousarray(probabilities, dtype=np.float16)
+        base_array = np.ascontiguousarray(encoded_bases, dtype=np.int8)
+        if probability_array.ndim != 1:
+            raise ValueError(
+                f"decoder task {task_index} probabilities must be 1D, "
+                f"got {probability_array.shape}"
+            )
+        if base_array.ndim != 1:
+            raise ValueError(
+                f"decoder task {task_index} encoded bases must be 1D, "
+                f"got {base_array.shape}"
+            )
+        if len(probability_array) != len(base_array):
+            raise ValueError(
+                f"decoder task {task_index} has {len(probability_array)} "
+                f"probabilities but {len(base_array)} encoded bases"
+            )
+        if len(probability_array) == 0:
+            raise ValueError(f"decoder task {task_index} is empty")
+        prepared_tasks.append((probability_array, base_array))
+    return prepared_tasks
+
+
+def prepare_transition_probabilities(transition_probabilities, prepared_tasks):
+    """Quantize one ``(bases, 4)`` transition-probability array per task for the decoders.
+
+    The columns follow ``TRANSITION_NAMES``: enter gene, leave gene, enter
+    intron, leave intron. Each value is the transition network's log-odds
+    against staying, in nats, aligned base by base with the task's
+    probabilities. The Viterbi implementations accumulate int32 fixed-point
+    scores, so the values come back rounded to ``FIXED_POINT_SCORE_SCALE``
+    units as contiguous int32 arrays.
+    """
+    transition_probabilities = list(transition_probabilities)
+    if len(transition_probabilities) != len(prepared_tasks):
+        raise ValueError(
+            f"{len(transition_probabilities)} transition-probability arrays for "
+            f"{len(prepared_tasks)} decoder tasks"
+        )
+    limits = np.iinfo(np.int32)
+    prepared_scores = []
+    for task_index, (log_odds, (probabilities, _)) in enumerate(
+        zip(transition_probabilities, prepared_tasks)
+    ):
+        scaled = np.rint(np.asarray(log_odds, dtype=np.float64) * FIXED_POINT_SCORE_SCALE)
+        if scaled.shape != (len(probabilities), TRANSITION_COUNT):
+            raise ValueError(
+                f"decoder task {task_index} transition probabilities must have shape "
+                f"{(len(probabilities), TRANSITION_COUNT)}, got {scaled.shape}"
+            )
+        if not np.all(np.isfinite(scaled)) or scaled.min() < limits.min or scaled.max() > limits.max:
+            raise OverflowError(f"decoder task {task_index} transition log-odds exceed the int32 fixed-point range")
+        prepared_scores.append(np.ascontiguousarray(scaled, dtype=np.int32))
+    return prepared_scores
+
+
+def build_cds_sequence(genome_sequence, transcript_order_exons, strand):
+    """Build a CDS string from genomic exon intervals in transcript order."""
+    if not transcript_order_exons:
+        return ""
+    if strand == "+":
+        return "".join(
+            genome_sequence[start:end + 1] for start, end in transcript_order_exons
+        )
+
+    genomic_order_sequence = "".join(
+        genome_sequence[start:end + 1]
+        for start, end in reversed(transcript_order_exons)
+    )
+    return reverse_complement(genomic_order_sequence)
+
+
+def compute_gff_cds_phases(transcript_order_exons, initial_cds_phase=0):
+    """Return GFF3 phases, preserving the first visible codon's phase."""
+    if initial_cds_phase not in (0, 1, 2):
+        raise ValueError("initial CDS phase must be 0, 1, or 2")
+
+    phases = []
+    consumed_bases = 0
+    for start, end in transcript_order_exons:
+        phase = (initial_cds_phase - consumed_bases) % 3
+        phases.append(phase)
+        consumed_bases += end - start + 1
+    return phases
+
+
+def has_internal_stop_codon(
+    cds_sequence,
+    translation_frame,
+    stop_codons=STOP_CODONS,
+):
+    """Return whether a translated CDS contains a non-terminal stop codon.
+
+    This is a diagnostic check: the decoder grammar should normally make the
+    result false. A stop is terminal only when its third base is the final CDS
+    base; a stop before an incomplete trailing codon is still internal.
+    """
+    if translation_frame not in (0, 1, 2):
+        raise ValueError("translation frame must be 0, 1, or 2")
+
+    codon_end = len(cds_sequence) - (
+        (len(cds_sequence) - translation_frame) % 3
+    )
+    for codon_start in range(translation_frame, codon_end, 3):
+        codon = cds_sequence[codon_start:codon_start + 3].upper()
+        if codon in stop_codons and codon_start + 3 < len(cds_sequence):
+            return True
+    return False
+
+
+def classify_cds_completion(
+    cds_sequence,
+    has_start_codon,
+    has_stop_codon,
+    has_short_intron=False,
+    stop_codons=STOP_CODONS,
+):
+    """Return ``(complete, reason)`` for a reconstructed CDS sequence.
+
+    Reasons are comma-separated so boundary-truncated predictions can state
+    every visible problem inside one GFF3 attribute, for example
+    ``no_start,no_stop``. ``short_intron``
+    records a decoder path with fewer than four bases between adjacent exons.
+    Start and stop presence comes from the decoded states, then the sequence
+    is checked as a consistency guard.
+    """
+    reasons = []
+    if len(cds_sequence) < 6:
+        reasons.append("too_short")
+    if len(cds_sequence) % 3 != 0:
+        reasons.append("not_multiple_of_3")
+    if has_short_intron:
+        reasons.append("short_intron")
+    if (
+        not has_start_codon
+        or len(cds_sequence) < 3
+        or cds_sequence[:3] not in START_CODONS
+    ):
+        reasons.append("no_start")
+    if (
+        not has_stop_codon
+        or len(cds_sequence) < 3
+        or cds_sequence[-3:] not in stop_codons
+    ):
+        reasons.append("no_stop")
+
+    if reasons:
+        # Comma-separated: a semicolon would end the GFF3 attribute.
+        return False, ",".join(reasons)
+    return True, None
+
+
+def codon_table_mismatch(
+    confidences, threshold=CONFIDENCE_THRESHOLD,
+    fraction=CODON_TABLE_WARNING_FRACTION, min_genes=CODON_TABLE_WARNING_MIN_GENES,
+):
+    """Return whether the confidences of the complete genes contradict the codon table.
+
+    A codon table that does not match the genome splits genes at codons the
+    genome reads through or fuses genes at stop codons it lacks, and the
+    confidence network gives such genes almost no confidence. The table is
+    contradicted when at least ``min_genes`` complete genes were scored and
+    fewer than ``fraction`` of them reach ``threshold``.
+    """
+    confidences = np.asarray(list(confidences), dtype=np.float64)
+    if len(confidences) < min_genes:
+        return False
+    return bool(np.mean(confidences >= threshold) < fraction)
+
+
+def has_short_internal_intron(exons):
+    """Return whether adjacent exons leave an internal intron under four bases."""
+    for exon_index in range(1, len(exons)):
+        previous_exon = exons[exon_index - 1]
+        next_exon = exons[exon_index]
+        if previous_exon[0] < next_exon[0]:
+            intron_length = next_exon[0] - previous_exon[1] - 1
+        else:
+            intron_length = previous_exon[0] - next_exon[1] - 1
+        if 0 < intron_length < MIN_INTERNAL_INTRON_LENGTH_BP:
+            return True
+    return False
+
+
+def translate_cds_sequence(
+    cds_sequence,
+    frame,
+    amino_acid_by_codon_index,
+):
+    """Translate one CDS sequence with a prepared codon-table lookup.
+
+    A single stop at the final complete codon is removed from the protein.
+    Any earlier stop indicates an invalid decoded CDS and raises an assertion
+    instead of silently truncating the output. An incomplete trailing codon is
+    ignored unless it follows a stop, in which case that stop is internal.
+    """
+    codon_end = len(cds_sequence) - ((len(cds_sequence) - frame) % 3)
+    if codon_end <= frame:
+        return ""
+
+    sequence_bytes = np.frombuffer(cds_sequence.encode("ascii"), dtype=np.uint8)
+    base_values = TRANSLATION_BASE_TO_VALUE[
+        sequence_bytes[frame:codon_end]
+    ].reshape(-1, 3)
+    codon_indices = (
+        base_values[:, 0] * TRANSLATION_BASE_COUNT * TRANSLATION_BASE_COUNT
+        + base_values[:, 1] * TRANSLATION_BASE_COUNT
+        + base_values[:, 2]
+    )
+    amino_acids = amino_acid_by_codon_index[codon_indices]
+
+    stop_positions = np.flatnonzero(amino_acids == ord("*"))
+    if stop_positions.size:
+        only_terminal_stop = (
+            stop_positions.size == 1
+            and stop_positions[0] == len(amino_acids) - 1
+            and codon_end == len(cds_sequence)
+        )
+        if not only_terminal_stop:
+            stop_index = int(stop_positions[0])
+            codon_start = frame + stop_index * 3
+            stop_codon = cds_sequence[codon_start:codon_start + 3].upper()
+            raise AssertionError(
+                "protein translation encountered an internal stop codon: "
+                f"codon={stop_codon}, cds_offset={codon_start}"
+            )
+        amino_acids = amino_acids[:-1]
+    return amino_acids.tobytes().decode("ascii")
+
+
+def assemble_gene_model(
+    record_id,
+    genome_sequence,
+    oriented_exons,
+    strand,
+    record_length,
+    first_cds_state,
+    last_cds_state,
+    stop_codons=STOP_CODONS,
+):
+    """Build one output gene from its exons and boundary CDS states."""
+    initial_cds_phase = initial_cds_phase_from_state(first_cds_state)
+    if strand == "+":
+        transcript_order_exons = [
+            (int(start), int(end)) for start, end in oriented_exons
+        ]
+    else:
+        transcript_order_exons = [
+            (record_length - 1 - int(end), record_length - 1 - int(start))
+            for start, end in oriented_exons
+        ]
+
+    cds_sequence = build_cds_sequence(genome_sequence, transcript_order_exons, strand)
+    if not cds_sequence:
+        return None
+
+    transcript_order_phases = compute_gff_cds_phases(
+        transcript_order_exons,
+        initial_cds_phase,
+    )
+    genomic_order = sorted(
+        zip(transcript_order_exons, transcript_order_phases),
+        key=lambda exon_and_phase: exon_and_phase[0][0],
+    )
+
+    assert not has_internal_stop_codon(
+        cds_sequence,
+        initial_cds_phase,
+        stop_codons,
+    ), (
+        "decoder produced an internal stop codon: "
+        f"record_id={record_id!r}, strand={strand}, "
+        f"phase={initial_cds_phase}, exons={transcript_order_exons}"
+    )
+    complete, incomplete_reason = classify_cds_completion(
+        cds_sequence,
+        first_cds_state == START_A,
+        last_cds_state == STOP,
+        has_short_intron=has_short_internal_intron(transcript_order_exons),
+        stop_codons=stop_codons,
+    )
+    return {
+        "record_id": record_id,
+        "strand": strand,
+        "cds": cds_sequence,
+        "translation_frame": initial_cds_phase,
+        "complete": complete,
+        "incomplete_reason": incomplete_reason,
+        "exons": [exon for exon, phase in genomic_order],
+        "phases": [phase for exon, phase in genomic_order],
+    }
+
+
+def format_gene_status(gene, separator):
+    """Format shared complete/incomplete metadata for GFF3 or FASTA."""
+    status = "complete" if gene["complete"] else "incomplete"
+    fields = [f"status={status}"]
+    if not gene["complete"]:
+        reason = gene.get("incomplete_reason") or "unknown"
+        fields.append(f"incomplete_reason={reason}")
+    return separator.join(fields)
+
+
+def format_gene_confidence(gene):
+    """Format the per-gene confidence for GFF3 attributes or FASTA headers.
+
+    Complete genes carry the confidence model's probability that the decoded
+    gene is exactly right. Incomplete genes are not scored and are written as
+    ``NA`` so that a confidence filter keeps them.
+    """
+    value = gene.get("confidence")
+    formatted_value = "NA" if value is None else f"{value:.6f}"
+    return f"confidence={formatted_value}"
+
+
+def write_gff3(genes, output_path, source="Carbon-A"):
+    """Write gene, mRNA, and CDS features in GFF3 format."""
+    with open(output_path, "w", encoding="utf-8", newline="\n") as handle:
+        handle.write("##gff-version 3\n")
+        for gene_number, gene in enumerate(genes, 1):
+            gene_id = f"gene{gene_number:06d}"
+            transcript_id = f"mrna{gene_number:06d}"
+            gene_start = gene["exons"][0][0] + 1
+            gene_end = gene["exons"][-1][1] + 1
+            status_attributes = format_gene_status(gene, separator=";")
+            status_attributes += ";" + format_gene_confidence(gene)
+
+            handle.write(
+                f"{gene['record_id']}\t{source}\tgene\t{gene_start}\t{gene_end}\t.\t"
+                f"{gene['strand']}\t.\tID={gene_id};{status_attributes}\n"
+            )
+            handle.write(
+                f"{gene['record_id']}\t{source}\tmRNA\t{gene_start}\t{gene_end}\t.\t"
+                f"{gene['strand']}\t.\tID={transcript_id};Parent={gene_id}\n"
+            )
+
+            for (start, end), phase in zip(gene["exons"], gene["phases"]):
+                handle.write(
+                    f"{gene['record_id']}\t{source}\tCDS\t{start + 1}\t{end + 1}\t.\t"
+                    f"{gene['strand']}\t{phase}\tID={transcript_id}.cds;Parent={transcript_id}\n"
+                )
+
+
+def write_cds_fasta(genes, output_path):
+    """Write predicted CDS sequences in FASTA format."""
+    with open(output_path, "w", encoding="utf-8", newline="\n") as handle:
+        for gene_number, gene in enumerate(genes, 1):
+            gene_start = gene["exons"][0][0] + 1
+            gene_end = gene["exons"][-1][1] + 1
+            header = (
+                f">gene{gene_number:06d} record_id={gene['record_id']} "
+                f"strand={gene['strand']} loc={gene_start}-{gene_end} "
+                f"phase={gene['translation_frame']} "
+                f"{format_gene_status(gene, separator=' ')}"
+            )
+            header += " " + format_gene_confidence(gene)
+
+            handle.write(header + "\n")
+            cds_sequence = gene["cds"]
+            for line_start in range(0, len(cds_sequence), 80):
+                handle.write(cds_sequence[line_start:line_start + 80] + "\n")
+
+
+def write_protein_fasta(
+    genes,
+    output_path,
+    codon_to_amino_acid,
+):
+    """Write translated protein sequences in FASTA format.
+
+    Returns the number of non-empty protein records written. Incomplete CDS
+    records keep their reason in the header so downstream filters can decide
+    whether to use them.
+    """
+    amino_acid_by_codon_index = build_translation_lookup(codon_to_amino_acid)
+    written_count = 0
+    with open(output_path, "w", encoding="utf-8", newline="\n") as handle:
+        for gene_number, gene in enumerate(genes, 1):
+            cds_sequence = gene["cds"]
+            protein_sequence = translate_cds_sequence(
+                cds_sequence,
+                gene["translation_frame"],
+                amino_acid_by_codon_index,
+            )
+            if not protein_sequence:
+                continue
+
+            gene_start = gene["exons"][0][0] + 1
+            gene_end = gene["exons"][-1][1] + 1
+            header = (
+                f">gene{gene_number:06d} record_id={gene['record_id']} "
+                f"strand={gene['strand']} loc={gene_start}-{gene_end} "
+                f"cds_len={len(cds_sequence)} aa_len={len(protein_sequence)} "
+                f"{format_gene_status(gene, separator=' ')}"
+            )
+            header += " " + format_gene_confidence(gene)
+
+            handle.write(header + "\n")
+            for line_start in range(0, len(protein_sequence), 80):
+                handle.write(protein_sequence[line_start:line_start + 80] + "\n")
+            written_count += 1
+
+    return written_count
+
+
+def load_nuclear_evaluation_accessions(genome_dir):
+    """Return the record accessions used for nuclear-gene evaluation.
+
+    Prefer assembled chromosomes. If none exist, use assembled linkage groups;
+    if those are also absent, use chromosome scaffolds. ``None`` means that no
+    sequence report is available for filtering.
+    """
+    sequence_report_path = Path(genome_dir) / "sequence_report.jsonl"
+    if not sequence_report_path.exists():
+        return None
+
+    assembled_chromosome_accessions = set()
+    assembled_linkage_group_accessions = set()
+    chromosome_scaffold_accessions = set()
+    with sequence_report_path.open(encoding="utf-8") as handle:
+        for line in handle:
+            line = line.strip()
+            if not line:
+                continue
+
+            row = json.loads(line)
+            accession = row.get("refseqAccession") or row.get("genbankAccession")
+            if not accession:
+                continue
+
+            assembly_unit = row.get("assemblyUnit")
+            location_type = row.get("assignedMoleculeLocationType")
+            if assembly_unit == NON_NUCLEAR_ASSEMBLY_UNIT:
+                continue
+
+            role = row.get("role")
+            if role == "assembled-molecule":
+                if location_type == "Chromosome":
+                    assembled_chromosome_accessions.add(accession)
+                elif location_type == "Linkage Group":
+                    assembled_linkage_group_accessions.add(accession)
+            elif location_type == "Chromosome" and (role or "").endswith("-scaffold"):
+                chromosome_scaffold_accessions.add(accession)
+
+    return (
+        assembled_chromosome_accessions
+        or assembled_linkage_group_accessions
+        or chromosome_scaffold_accessions
+    )

annotation_pipeline/src/viterbi.py

@@ -0,0 +1,940 @@
+"""Viterbi algorithm: the highest-scoring gene labeling of one strand.
+
+The algorithm labels every base of a strand, in transcript orientation, as
+intergenic, exon, or intron. The score of a labeling is the sum over bases of
+log P(CDS) for exon bases and log(1 - P(CDS)) for all other bases, plus the
+log-probability of the transition the transition network assigns at every
+base where a segment begins or a gene closes, plus a fixed log-prior for each
+intron's splice class. The implementation adds each transition's log-odds
+against staying instead of its log-probability; the two differ by the sum of
+log P(stay) over all boundaries, which is the same for every labeling. The gene
+grammar decides which labelings are admissible: a gene opens at an ATG start
+codon and closes at the first in-frame stop codon of the selected genetic
+code, an intron opens at a donor site and closes at an acceptor site of a
+permitted class, and the reading frame continues across each intron. The
+implementation expands the three labels into 30 states so that every
+condition can be checked from local context.
+
+This module holds the readable Numba reference implementation, the
+interpretation of a decoded state path as genes and exons, and the choice of
+Viterbi backend: the equivalent native implementations in viterbi_backends/
+(C++/OpenMP and CUDA) return identical paths far faster.
+
+State design:
+
+    To recognize stop codons, the decoder only needs to remember which stop
+    prefix, if any, has already been seen inside the current codon. That memory
+    is stored as one of six prefix types. Codon phase here means how many CDS
+    bases, 0, 1, or 2, have been consumed in the current codon:
+
+        prefix type 0: codon phase 0; codon boundary, no open prefix
+        prefix type 1: codon phase 1; one-base prefix T
+        prefix type 2: codon phase 1; one-base prefix other
+        prefix type 3: codon phase 2; prefix TA, so A/G completes TAA/TAG
+        prefix type 4: codon phase 2; prefix TG, so A completes TGA
+        prefix type 5: codon phase 2; other prefix, so no stop can be completed
+
+    GFF3 phase uses a different convention: it records how many visible bases
+    to skip before the next complete codon, rather than bases already consumed.
+
+    Exon states carry the current prefix type while emitting coding bases.
+    Intron states carry the prefix type from the exon before the donor, so the
+    acceptor resumes CDS in the same frame. GT-AG and GC-AG share one intron
+    state table because both exit at AG, but each uses its own intron prior.
+    AT-AC uses a separate state table because it exits at AC.
+
+    The active state ids are defined in utils.py:
+
+    0   INTERGENIC                              non-coding outside genes
+    1   START_A                                 first base of start codon ATG
+    2   START_AT                                first two bases of start codon ATG
+    3   EXON_PREFIX_TYPE_0                      CDS, carries prefix type 0
+    4   EXON_PREFIX_TYPE_1                      CDS, carries prefix type 1
+    5   EXON_PREFIX_TYPE_2                      CDS, carries prefix type 2
+    6   EXON_PREFIX_TYPE_3                      CDS, carries prefix type 3
+    7   EXON_PREFIX_TYPE_4                      CDS, carries prefix type 4
+    8   EXON_PREFIX_TYPE_5                      CDS, carries prefix type 5
+    9   STOP                                    terminal in-frame stop codon
+    10  GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0     first GT/GC-AG intron base, carries prefix type 0
+    11  GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_1     first GT/GC-AG intron base, carries prefix type 1
+    12  GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_2     first GT/GC-AG intron base, carries prefix type 2
+    13  GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_3     first GT/GC-AG intron base, carries prefix type 3
+    14  GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_4     first GT/GC-AG intron base, carries prefix type 4
+    15  GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_5     first GT/GC-AG intron base, carries prefix type 5
+    16  GT_GC_AG_INTRON_READY_PREFIX_TYPE_0     later GT/GC-AG intron base, carries prefix type 0
+    17  GT_GC_AG_INTRON_READY_PREFIX_TYPE_1     later GT/GC-AG intron base, carries prefix type 1
+    18  GT_GC_AG_INTRON_READY_PREFIX_TYPE_2     later GT/GC-AG intron base, carries prefix type 2
+    19  GT_GC_AG_INTRON_READY_PREFIX_TYPE_3     later GT/GC-AG intron base, carries prefix type 3
+    20  GT_GC_AG_INTRON_READY_PREFIX_TYPE_4     later GT/GC-AG intron base, carries prefix type 4
+    21  GT_GC_AG_INTRON_READY_PREFIX_TYPE_5     later GT/GC-AG intron base, carries prefix type 5
+    22  AT_AC_INTRON_PREFIX_TYPE_0              AT-AC intron, carries prefix type 0
+    23  AT_AC_INTRON_PREFIX_TYPE_1              AT-AC intron, carries prefix type 1
+    24  AT_AC_INTRON_PREFIX_TYPE_2              AT-AC intron, carries prefix type 2
+    25  AT_AC_INTRON_PREFIX_TYPE_3              AT-AC intron, carries prefix type 3
+    26  AT_AC_INTRON_PREFIX_TYPE_4              AT-AC intron, carries prefix type 4
+    27  AT_AC_INTRON_PREFIX_TYPE_5              AT-AC intron, carries prefix type 5
+    28  LEADING_BOUNDARY_INTRON                 intron extending from the left record boundary
+    29  TRAILING_BOUNDARY_INTRON                intron extending to the right record boundary
+
+Transition rules:
+
+    INTERGENIC self-loops or starts a start codon ATG. START_A -> START_AT ->
+    exon prefix type 0. Exons consume coding bases and update prefix type;
+    selected stops enter STOP. STOP exits to INTERGENIC. Exons may enter
+    GT-AG/GC-AG introns at G followed by T/C, or AT-AC introns at A followed
+    by T. GT/GC-AG introns enter FIRST, then must consume one more intron base
+    before becoming READY; only READY can exit after AG. This prevents exon A
+    plus intron G in an AGT/AGC boundary from being misread as an acceptor.
+    AT-AC can use one state table because AT cannot overlap AC in the same way.
+
+    Opening an intron adds the explicit log prior for its splice class. The
+    supported frequencies are 0.990 for GT-AG, 0.009 for GC-AG, and 0.001 for
+    AT-AC, each divided by the 0.99 GT-AG baseline before taking the logarithm.
+    The GT-AG prior is therefore exactly zero, but remains an explicit decoder
+    parameter in the Numba, CPU, and CUDA implementations.
+
+    Except for the two boundary-intron guards described below, a sequence base
+    N is a free wildcard. When a transition needs its identity, N may satisfy
+    any A, T, C, or G requirement without changing the score. Coding N bases
+    branch into every compatible prefix state, so subsequent stop-codon
+    decisions remain consistent with the selected branch.
+
+    Splice-signal matching deliberately does not commit N to one concrete base.
+    The same N can therefore satisfy overlapping donor and acceptor checks with
+    different interpretations, allowing an internal intron shorter than the
+    four bases needed to contain both motifs. No such intron occurred in the
+    26-species validation set. Such a path must be supported by the transition
+    scores at both ends of the intron and still pays the splice-class prior. If
+    it wins, the model strongly prefers those bases to be non-coding, while the
+    ambiguous sequence provides too little evidence for a reliable repair. The path is retained and its gene is reported as
+    incomplete with reason short_intron instead of expanding the decoder state
+    space solely for this rare condition.
+
+    The transition network's four probabilities enter the score as log-odds
+    against staying, charged on the base where the new segment begins or, for
+    leaving a gene, on the base that closes it:
+
+        INTERGENIC -> START_A                 enter gene    first start-codon base
+        any state -> STOP                     leave gene    last stop-codon base
+        exon -> ordinary/trailing intron      enter intron  first intron base
+        ordinary/leading intron -> exon       leave intron  first exon base after the intron
+
+    Every other state change adds nothing. The first base is initialized
+    the same way: a start codon beginning there is charged enter gene and a
+    stop codon completed there is charged leave gene, mirroring the right
+    edge, where a stop codon completed at the last base is charged leave gene.
+    When an intron exit completes a stop codon that the intron split, the base
+    closes the gene and is charged leave gene.
+
+    The two boundary-intron states represent genes truncated inside an intron.
+    LEADING_BOUNDARY_INTRON is the only intron state allowed at the first base.
+    Any prefix type may be chosen when LEADING_BOUNDARY_INTRON exits, so the
+    state does not need to carry one while it remains intronic. If a concrete
+    AG or AC acceptor is immediately followed by a concrete ATG, the boundary
+    state may remain intronic but cannot exit there. Any path that treats this
+    ATG as CDS must therefore use the normal start-codon states instead of
+    borrowing an unknown phase.
+
+    TRAILING_BOUNDARY_INTRON can be entered from any exon at a legal donor and
+    can only self-loop to the record end. Entry is forbidden when the three
+    concrete bases immediately before a concrete donor form an enabled stop
+    codon, regardless of the path's carried prefix type. This prevents a path
+    from choosing another phase merely to hide a terminal stop. These two
+    boundary-only guards require every base in the acceptor/codon or
+    codon/donor pattern to be present and different from N. If any base is
+    missing or N, the hard guard does not apply and model scores choose between
+    the otherwise legal paths. N remains a wildcard for the underlying splice
+    and ordinary start/stop transitions. An ordinary intron must be entered
+    from an exon at a visible donor and must exit to an exon at a visible
+    acceptor. Only a boundary-intron state may represent an intron that crosses
+    a record edge.
+
+    max_boundary_intron_length limits how many consecutive bases either
+    boundary-intron state may emit. The leading state can only extend that far
+    from the left edge; the trailing state can only be entered within that
+    distance of the right edge. The production default is 20,000 bp. Zero
+    disables both boundary-intron states, so incomplete genes can then only arise
+    when a record edge truncates an exon. A value of -1 leaves boundary-intron
+    lengths unlimited. Exon-truncated incomplete genes remain legal in every mode.
+
+Record boundaries:
+
+    A record may begin or end inside a gene. Missing motif bases outside the
+    record are treated as unknown: they may complete a boundary-spanning
+    start, stop, donor, or acceptor. This relaxation applies only where context
+    is unavailable; every internal transition still requires its full motif.
+"""
+
+import importlib
+import os
+import sys
+import warnings
+
+import numpy as np
+from numba import njit
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from constants import (
+    AT_AC_INTRON_LOG_PRIOR,
+    BASE_A,
+    BASE_C,
+    BASE_G,
+    BASE_N,
+    BASE_T,
+    ENTER_GENE,
+    ENTER_INTRON,
+    FIXED_POINT_SCORE_SCALE,
+    GC_AG_INTRON_LOG_PRIOR,
+    GT_AG_INTRON_LOG_PRIOR,
+    LEAVE_GENE,
+    LEAVE_INTRON,
+    MAX_BOUNDARY_INTRON_LENGTH_BP,
+    MAX_PROB,
+    MIN_PROB,
+)
+from utils import (
+    PREFIX_TYPE_0, PREFIX_TYPE_1, PREFIX_TYPE_2,
+    PREFIX_TYPE_3, PREFIX_TYPE_4, PREFIX_TYPE_5, PREFIX_TYPE_COUNT,
+    STATE_COUNT, INTERGENIC, START_A, START_AT, EXON_PREFIX_TYPE_0, STOP,
+    LEADING_BOUNDARY_INTRON, TRAILING_BOUNDARY_INTRON,
+    EXON_STATE_BY_PREFIX_TYPE,
+    GT_GC_AG_INTRON_FIRST_STATE_BY_PREFIX_TYPE,
+    GT_GC_AG_INTRON_READY_STATE_BY_PREFIX_TYPE,
+    AT_AC_INTRON_STATE_BY_PREFIX_TYPE,
+    prepare_decode_tasks, prepare_transition_probabilities,
+    validate_max_boundary_intron_length,
+)
+
+
+@njit(cache=True, inline="always")
+def keep_best_transition(next_scores, predecessors, target_state, candidate_score, source_state):
+    """Apply one Viterbi max update for a candidate transition."""
+    # Standard Viterbi max update: keep the best score and remember where it
+    # came from so backtrace can reconstruct the path.
+    if candidate_score > next_scores[target_state]:
+        next_scores[target_state] = candidate_score
+        predecessors[target_state] = source_state
+
+
+@njit(cache=True, inline="always")
+def base_matches(observed_base, expected_base):
+    """Return whether a concrete base or the N wildcard matches a base."""
+    return observed_base == expected_base or observed_base == BASE_N
+
+
+@njit(cache=True, inline="always")
+def is_exact_enabled_stop_codon(
+    first_base, second_base, third_base,
+    stop_taa, stop_tag, stop_tga,
+):
+    """Return whether three concrete bases are an enabled stop codon."""
+    if first_base != BASE_T:
+        return False
+    if second_base == BASE_A:
+        if third_base == BASE_A and stop_taa:
+            return True
+        if third_base == BASE_G and stop_tag:
+            return True
+    return second_base == BASE_G and third_base == BASE_A and stop_tga
+
+
+@njit(cache=True, inline="always")
+def next_coding_state_after_base(prefix_type, base, stop_taa, stop_tag, stop_tga):
+    """Return the exon/STOP state after emitting one coding base."""
+    # Prefix type 0 means the previous coding base ended a codon. A new T can
+    # still become TAA/TAG/TGA; every other base cannot start a stop codon.
+    if prefix_type == PREFIX_TYPE_0:
+        if base == BASE_T:
+            return EXON_STATE_BY_PREFIX_TYPE[PREFIX_TYPE_1]
+        return EXON_STATE_BY_PREFIX_TYPE[PREFIX_TYPE_2]
+
+    # Prefix type 1 means the current codon prefix is exactly T.
+    if prefix_type == PREFIX_TYPE_1:
+        if base == BASE_A:
+            return EXON_STATE_BY_PREFIX_TYPE[PREFIX_TYPE_3]
+        if base == BASE_G:
+            return EXON_STATE_BY_PREFIX_TYPE[PREFIX_TYPE_4]
+        return EXON_STATE_BY_PREFIX_TYPE[PREFIX_TYPE_5]
+
+    # Prefix type 2 is a one-base prefix other than T, so the second base cannot
+    # create a possible standard stop prefix.
+    if prefix_type == PREFIX_TYPE_2:
+        return EXON_STATE_BY_PREFIX_TYPE[PREFIX_TYPE_5]
+
+    # Prefix type 3 is TA. A or G completes TAA/TAG if those stops are enabled.
+    if prefix_type == PREFIX_TYPE_3:
+        if base == BASE_A and stop_taa:
+            return STOP
+        if base == BASE_G and stop_tag:
+            return STOP
+        return EXON_PREFIX_TYPE_0
+
+    # Prefix type 4 is TG. A completes TGA if that stop is enabled.
+    if prefix_type == PREFIX_TYPE_4:
+        if base == BASE_A and stop_tga:
+            return STOP
+        return EXON_PREFIX_TYPE_0
+
+    return EXON_PREFIX_TYPE_0
+
+
+@njit(cache=True, inline="always")
+def emit_coding_base(
+    current_scores, next_scores, predecessors,
+    source_state, prefix_type, coding_score,
+    leave_intron_log_odds, leave_gene_log_odds,
+    base, stop_taa, stop_tag, stop_tga,
+):
+    """Emit the current base as CDS from an exon or intron-acceptor state.
+
+    ``leave_intron_log_odds`` is the log-odds of returning to an exon;
+    it is zero when the source state is already an exon. A base that completes
+    an enabled stop codon closes the gene and is charged ``leave_gene_log_odds``
+    instead.
+    """
+    # Used both by normal exon continuation and by intron acceptor exits. In
+    # both cases the current base is emitted as CDS and updates prefix type.
+    source_score = current_scores[source_state]
+    if source_score == -np.inf:
+        return
+
+    base_score = source_score + coding_score
+    if base == BASE_N:
+        # Resolve N independently along every legal A/T/C/G branch. The
+        # selected prefix state remembers the consequence of that choice.
+        for resolved_base in (BASE_A, BASE_T, BASE_C, BASE_G):
+            target_state = next_coding_state_after_base(
+                prefix_type, resolved_base, stop_taa, stop_tag, stop_tga
+            )
+            transition_log_odds = leave_gene_log_odds if target_state == STOP else leave_intron_log_odds
+            keep_best_transition(
+                next_scores, predecessors, target_state,
+                base_score + transition_log_odds, source_state,
+            )
+    else:
+        target_state = next_coding_state_after_base(
+            prefix_type, base, stop_taa, stop_tag, stop_tga
+        )
+        transition_log_odds = leave_gene_log_odds if target_state == STOP else leave_intron_log_odds
+        keep_best_transition(
+            next_scores, predecessors, target_state,
+            base_score + transition_log_odds, source_state,
+        )
+
+
+@njit(cache=True)
+def initialize_first_base_scores(
+    coding_scores, noncoding_scores, transition_probabilities, bases, record_length,
+    stop_taa, stop_tag, stop_tga, boundary_introns_enabled,
+    current_scores,
+):
+    """Initialize first-position state scores, allowing truncated genes at record edges.
+
+    Transitions across the left record edge are charged like interior ones:
+    a start codon beginning at the first base enters the gene at boundary 0
+    (``enter_gene[0]``), and a stop codon completed at the first base leaves
+    it at boundary 1 (``leave_gene[0]``). Every other first-base state stays
+    in whatever the record began with and is charged nothing.
+    """
+    # The first base is not forced to be intergenic. A record may start in the
+    # middle of a gene or intron, so initialize every start state that is
+    # compatible with this observed base and an unknown left context.
+    # That is why several exon prefix types and the leading boundary-intron
+    # state can be reachable at position 0.
+    for state in range(STATE_COUNT):
+        # -np.inf means this state has no legal path ending here.
+        current_scores[state] = -np.inf
+
+    base = bases[0]
+    coding_score = coding_scores[0]
+    noncoding_score = noncoding_scores[0]
+
+    current_scores[INTERGENIC] = noncoding_score
+
+    # These exon initial states enumerate all prefix types compatible with one
+    # observed base and unknown bases to the left of the record boundary. They
+    # remain legal when boundary introns are disabled, allowing a record to
+    # begin inside an exon.
+    current_scores[EXON_STATE_BY_PREFIX_TYPE[PREFIX_TYPE_0]] = coding_score
+    current_scores[EXON_STATE_BY_PREFIX_TYPE[PREFIX_TYPE_5]] = coding_score
+    if base == BASE_N:
+        for prefix_type in range(PREFIX_TYPE_COUNT):
+            current_scores[EXON_STATE_BY_PREFIX_TYPE[prefix_type]] = coding_score
+    else:
+        one_base_prefix_type = PREFIX_TYPE_1 if base == BASE_T else PREFIX_TYPE_2
+        current_scores[EXON_STATE_BY_PREFIX_TYPE[one_base_prefix_type]] = coding_score
+        if base == BASE_A:
+            current_scores[EXON_STATE_BY_PREFIX_TYPE[PREFIX_TYPE_3]] = coding_score
+        elif base == BASE_G:
+            current_scores[EXON_STATE_BY_PREFIX_TYPE[PREFIX_TYPE_4]] = coding_score
+
+    # The record may begin with the last observed base of a stop codon whose
+    # first two bases lie outside the left boundary.
+    can_end_stop = (
+        base_matches(base, BASE_A) and (stop_taa or stop_tga)
+    ) or (base_matches(base, BASE_G) and stop_tag)
+    if can_end_stop:
+        current_scores[STOP] = coding_score + transition_probabilities[0, LEAVE_GENE]
+
+    if boundary_introns_enabled:
+        # A record that starts inside an intron has no observable preceding
+        # exon, donor class, or codon phase. Ordinary intron states remain
+        # reserved for visible donors.
+        current_scores[LEADING_BOUNDARY_INTRON] = noncoding_score
+
+    # Missing right-edge bases are unknown, so an ATG truncated by the record
+    # boundary remains legal.
+    second_start_base_matches = record_length == 1 or base_matches(bases[1], BASE_T)
+    third_start_base_matches = record_length <= 2 or base_matches(bases[2], BASE_G)
+    has_atg_start = (
+        base_matches(base, BASE_A)
+        and second_start_base_matches
+        and third_start_base_matches
+    )
+    if has_atg_start:
+        current_scores[START_A] = coding_score + transition_probabilities[0, ENTER_GENE]
+
+
+@njit(cache=True)
+def run_viterbi_forward(
+    coding_scores, noncoding_scores, transition_probabilities, bases, record_length,
+    stop_taa, stop_tag, stop_tga,
+    gt_ag_intron_log_prior, gc_ag_intron_log_prior, at_ac_intron_log_prior,
+    max_boundary_intron_length,
+    current_scores, backpointers,
+):
+    """Run the source-state Viterbi forward pass and fill backpointers.
+
+    ``transition_probabilities`` is the ``(record_length, 4)`` array of transition
+    network scores in nats, in ``TRANSITION_NAMES`` column order.
+    """
+    initialize_first_base_scores(
+        coding_scores, noncoding_scores, transition_probabilities, bases, record_length,
+        stop_taa, stop_tag, stop_tga, max_boundary_intron_length != 0,
+        current_scores,
+    )
+    next_scores = np.empty(STATE_COUNT, np.float64)
+    predecessors = np.empty(STATE_COUNT, np.int8)
+
+    for position in range(1, record_length):
+        base = bases[position]
+        coding_score = coding_scores[position]
+        noncoding_score = noncoding_scores[position]
+        enter_gene_log_odds = transition_probabilities[position, ENTER_GENE]
+        leave_gene_log_odds = transition_probabilities[position, LEAVE_GENE]
+        enter_intron_log_odds = transition_probabilities[position, ENTER_INTRON]
+        leave_intron_log_odds = transition_probabilities[position, LEAVE_INTRON]
+
+        # Donor tests look at the current and next base. At the final base the
+        # unavailable next base may complete a donor outside the record.
+        next_base_is_missing = position + 1 == record_length
+        next_base = -1 if next_base_is_missing else bases[position + 1]
+        has_gt_ag_donor = base_matches(base, BASE_G) and (
+            next_base_is_missing or base_matches(next_base, BASE_T)
+        )
+        has_gc_ag_donor = base_matches(base, BASE_G) and (
+            next_base_is_missing or base_matches(next_base, BASE_C)
+        )
+
+        # Acceptor tests look behind the current base. At position 1, an A just
+        # outside the record may combine with position 0 G/C to complete AG/AC.
+        acceptor_first_base_matches = position == 1 or base_matches(bases[position - 2], BASE_A)
+        has_ag_acceptor = acceptor_first_base_matches and base_matches(bases[position - 1], BASE_G)
+        has_ac_acceptor = acceptor_first_base_matches and base_matches(bases[position - 1], BASE_C)
+
+        third_start_base_matches = position + 2 >= record_length or base_matches(
+            bases[position + 2], BASE_G
+        )
+        has_atg_start = (
+            base_matches(base, BASE_A)
+            and (next_base_is_missing or base_matches(next_base, BASE_T))
+            and third_start_base_matches
+        )
+        has_exact_acceptor_followed_by_atg = (
+            position >= 2
+            and position + 2 < record_length
+            and bases[position - 2] == BASE_A
+            and (
+                bases[position - 1] == BASE_G
+                or bases[position - 1] == BASE_C
+            )
+            and base == BASE_A
+            and next_base == BASE_T
+            and bases[position + 2] == BASE_G
+        )
+        has_at_ac_donor = base_matches(base, BASE_A) and (
+            next_base_is_missing or base_matches(next_base, BASE_T)
+        )
+
+        # Boundary-intron guards require exact five-base patterns. N and
+        # missing record-edge context still participate in ordinary motif
+        # matching, but never provide enough evidence for either hard guard.
+        has_exact_enabled_stop_followed_by_donor = False
+        has_any_donor = has_gt_ag_donor or has_gc_ag_donor or has_at_ac_donor
+        has_exact_donor = (
+            position + 1 < record_length
+            and (
+                (base == BASE_G and (next_base == BASE_T or next_base == BASE_C))
+                or (base == BASE_A and next_base == BASE_T)
+            )
+        )
+        if has_exact_donor and position >= 3:
+            has_exact_enabled_stop_followed_by_donor = is_exact_enabled_stop_codon(
+                bases[position - 3], bases[position - 2], bases[position - 1],
+                stop_taa, stop_tag, stop_tga,
+            )
+
+        # A leading self-loop here makes its run position + 1 bases long.
+        # Entering the trailing state here emits record_length - position bases
+        # through the right edge. No boundary counter state is needed.
+        boundary_introns_enabled = max_boundary_intron_length != 0
+        boundary_intron_length_is_unlimited = max_boundary_intron_length == -1
+        can_extend_leading_boundary_intron = (
+            boundary_introns_enabled
+            and (
+                boundary_intron_length_is_unlimited
+                or position < max_boundary_intron_length
+            )
+        )
+        can_enter_trailing_boundary_intron = (
+            boundary_introns_enabled
+            and (
+                boundary_intron_length_is_unlimited
+                or record_length - position <= max_boundary_intron_length
+            )
+        )
+
+        for state in range(STATE_COUNT):
+            # Reset the next column to unreachable before source states brush
+            # their legal transitions into it.
+            next_scores[state] = -np.inf
+            predecessors[state] = 0
+
+        # 1. Intergenic stays intergenic or starts a start codon ATG.
+        if current_scores[INTERGENIC] != -np.inf:
+            keep_best_transition(
+                next_scores, predecessors, INTERGENIC,
+                current_scores[INTERGENIC] + noncoding_score,
+                INTERGENIC,
+            )
+            if has_atg_start:
+                keep_best_transition(
+                    next_scores, predecessors, START_A,
+                    current_scores[INTERGENIC] + coding_score + enter_gene_log_odds,
+                    INTERGENIC,
+                )
+
+        # 2. Stop exits to intergenic.
+        if current_scores[STOP] != -np.inf:
+            keep_best_transition(
+                next_scores, predecessors, INTERGENIC,
+                current_scores[STOP] + noncoding_score,
+                STOP,
+            )
+
+        # 3. Start codon ATG consumes T and G, then enters exon prefix type 0.
+        if current_scores[START_A] != -np.inf:
+            keep_best_transition(
+                next_scores, predecessors, START_AT,
+                current_scores[START_A] + coding_score, START_A,
+            )
+        if current_scores[START_AT] != -np.inf:
+            keep_best_transition(
+                next_scores, predecessors, EXON_PREFIX_TYPE_0,
+                current_scores[START_AT] + coding_score,
+                START_AT,
+            )
+
+        # 4. A leading boundary intron self-loops or exits at either acceptor
+        # class. It carries no prefix type because any prefix type may be chosen
+        # when it exits. It cannot exit where an acceptor is immediately followed
+        # by an exact, fully visible ATG; that ATG must use the normal
+        # start-codon states. Any N in the acceptor or ATG leaves both paths
+        # available for the model scores to resolve.
+        if current_scores[LEADING_BOUNDARY_INTRON] != -np.inf:
+            if can_extend_leading_boundary_intron:
+                keep_best_transition(
+                    next_scores, predecessors,
+                    LEADING_BOUNDARY_INTRON,
+                    current_scores[LEADING_BOUNDARY_INTRON] + noncoding_score,
+                    LEADING_BOUNDARY_INTRON,
+                )
+            can_exit_leading_intron = (
+                (has_ag_acceptor or has_ac_acceptor)
+                and not has_exact_acceptor_followed_by_atg
+            )
+            if can_exit_leading_intron:
+                for prefix_type in range(PREFIX_TYPE_COUNT):
+                    emit_coding_base(
+                        current_scores, next_scores, predecessors,
+                        LEADING_BOUNDARY_INTRON, prefix_type,
+                        coding_score, leave_intron_log_odds, leave_gene_log_odds,
+                        base, stop_taa, stop_tag, stop_tga,
+                    )
+
+        # 5. Exons emit coding bases, open an ordinary intron that must later
+        # return to an exon, or commit to an intron that reaches the right edge.
+        for prefix_type in range(PREFIX_TYPE_COUNT):
+            exon_state = EXON_STATE_BY_PREFIX_TYPE[prefix_type]
+            exon_score = current_scores[exon_state]
+            if exon_score == -np.inf:
+                continue
+
+            # Continue the exon by emitting this base as coding. This may enter
+            # STOP when the prefix type plus current base forms an enabled stop.
+            emit_coding_base(
+                current_scores, next_scores, predecessors,
+                exon_state, prefix_type,
+                coding_score, 0.0, leave_gene_log_odds,
+                base, stop_taa, stop_tag, stop_tga,
+            )
+
+            # GT-AG and GC-AG share states because both wait for AG after entry.
+            # Each class uses its explicit log prior; the production GT-AG
+            # prior happens to be zero. Entry targets FIRST so the next base
+            # cannot immediately be treated as an exon base after a fake AG
+            # acceptor spanning the exon/intron boundary.
+            if has_gt_ag_donor:
+                keep_best_transition(
+                    next_scores, predecessors,
+                    GT_GC_AG_INTRON_FIRST_STATE_BY_PREFIX_TYPE[prefix_type],
+                    exon_score + noncoding_score + gt_ag_intron_log_prior + enter_intron_log_odds,
+                    exon_state,
+                )
+            if has_gc_ag_donor:
+                keep_best_transition(
+                    next_scores, predecessors,
+                    GT_GC_AG_INTRON_FIRST_STATE_BY_PREFIX_TYPE[prefix_type],
+                    exon_score + noncoding_score + gc_ag_intron_log_prior + enter_intron_log_odds,
+                    exon_state,
+                )
+
+            # U12 AT-AC introns are represented separately because they exit at
+            # AC instead of AG.
+            if has_at_ac_donor:
+                keep_best_transition(
+                    next_scores, predecessors,
+                    AT_AC_INTRON_STATE_BY_PREFIX_TYPE[prefix_type],
+                    exon_score + noncoding_score + at_ac_intron_log_prior + enter_intron_log_odds,
+                    exon_state,
+                )
+
+            if (
+                has_any_donor
+                and not has_exact_enabled_stop_followed_by_donor
+                and can_enter_trailing_boundary_intron
+            ):
+                if has_gt_ag_donor:
+                    keep_best_transition(
+                        next_scores, predecessors,
+                        TRAILING_BOUNDARY_INTRON,
+                        exon_score + noncoding_score + gt_ag_intron_log_prior + enter_intron_log_odds,
+                        exon_state,
+                    )
+                if has_gc_ag_donor:
+                    keep_best_transition(
+                        next_scores, predecessors,
+                        TRAILING_BOUNDARY_INTRON,
+                        exon_score + noncoding_score + gc_ag_intron_log_prior + enter_intron_log_odds,
+                        exon_state,
+                    )
+                if has_at_ac_donor:
+                    keep_best_transition(
+                        next_scores, predecessors,
+                        TRAILING_BOUNDARY_INTRON,
+                        exon_score + noncoding_score + at_ac_intron_log_prior + enter_intron_log_odds,
+                        exon_state,
+                    )
+
+        # 6. GT-AG/GC-AG FIRST states must consume one more intron base before
+        # acceptor exit is allowed.
+        for prefix_type in range(PREFIX_TYPE_COUNT):
+            intron_state = GT_GC_AG_INTRON_FIRST_STATE_BY_PREFIX_TYPE[prefix_type]
+            if current_scores[intron_state] == -np.inf:
+                continue
+            keep_best_transition(
+                next_scores, predecessors,
+                GT_GC_AG_INTRON_READY_STATE_BY_PREFIX_TYPE[prefix_type],
+                current_scores[intron_state] + noncoding_score,
+                intron_state,
+            )
+
+        # 7. GT-AG/GC-AG READY states self-loop or exit at AG.
+        for prefix_type in range(PREFIX_TYPE_COUNT):
+            intron_state = GT_GC_AG_INTRON_READY_STATE_BY_PREFIX_TYPE[prefix_type]
+            if current_scores[intron_state] == -np.inf:
+                continue
+            # Staying in an intron emits the current base as non-coding.
+            keep_best_transition(
+                next_scores, predecessors, intron_state,
+                current_scores[intron_state] + noncoding_score,
+                intron_state,
+            )
+            if has_ag_acceptor:
+                # Exiting at AG emits the current base as CDS and resumes with
+                # the carried prefix type.
+                emit_coding_base(
+                    current_scores, next_scores, predecessors,
+                    intron_state, prefix_type,
+                    coding_score, leave_intron_log_odds, leave_gene_log_odds,
+                    base, stop_taa, stop_tag, stop_tga,
+                )
+
+        # 8. AT-AC introns self-loop or exit at AC. AT cannot overlap AC the
+        # way GT/GC can overlap AG, so one state table is enough.
+        for prefix_type in range(PREFIX_TYPE_COUNT):
+            intron_state = AT_AC_INTRON_STATE_BY_PREFIX_TYPE[prefix_type]
+            if current_scores[intron_state] == -np.inf:
+                continue
+            # Same as GT/GC-AG introns, but the acceptor motif is AC.
+            keep_best_transition(
+                next_scores, predecessors, intron_state,
+                current_scores[intron_state] + noncoding_score,
+                intron_state,
+            )
+            if has_ac_acceptor:
+                emit_coding_base(
+                    current_scores, next_scores, predecessors,
+                    intron_state, prefix_type,
+                    coding_score, leave_intron_log_odds, leave_gene_log_odds,
+                    base, stop_taa, stop_tag, stop_tga,
+                )
+
+        # 9. A trailing boundary intron never exits back to an exon.
+        if current_scores[TRAILING_BOUNDARY_INTRON] != -np.inf:
+            keep_best_transition(
+                next_scores, predecessors,
+                TRAILING_BOUNDARY_INTRON,
+                current_scores[TRAILING_BOUNDARY_INTRON] + noncoding_score,
+                TRAILING_BOUNDARY_INTRON,
+            )
+
+        for state in range(STATE_COUNT):
+            current_scores[state] = next_scores[state]
+            backpointers[position - 1, state] = predecessors[state]
+
+
+@njit(cache=True)
+def best_final_state(current_scores):
+    """Return the best state that may legally end a record.
+
+    Ordinary introns must return to an exon. Only TRAILING_BOUNDARY_INTRON may
+    represent an intron that continues beyond the right record boundary. Exon
+    states may end at the boundary to represent an exon-truncated gene. Thus,
+    when boundary introns are disabled, exon truncation is the only way a
+    decoded gene can remain incomplete.
+    """
+    best_state = INTERGENIC
+    for state in range(START_A, STOP + 1):
+        if current_scores[state] > current_scores[best_state]:
+            best_state = state
+    if current_scores[TRAILING_BOUNDARY_INTRON] > current_scores[best_state]:
+        best_state = TRAILING_BOUNDARY_INTRON
+    return best_state
+
+
+@njit(cache=True, nogil=True)
+def decode_one_track(
+    coding_scores, noncoding_scores, transition_probabilities, bases, record_length,
+    stop_taa, stop_tag, stop_tga,
+    gt_ag_intron_log_prior, gc_ag_intron_log_prior, at_ac_intron_log_prior,
+    max_boundary_intron_length,
+    path,
+):
+    """Decode one oriented probability track into a per-base state path."""
+    current_scores = np.empty(STATE_COUNT, np.float64)
+    backpointers = np.empty((record_length, STATE_COUNT), np.int8)
+    run_viterbi_forward(
+        coding_scores, noncoding_scores, transition_probabilities, bases, record_length,
+        stop_taa, stop_tag, stop_tga,
+        gt_ag_intron_log_prior, gc_ag_intron_log_prior, at_ac_intron_log_prior,
+        max_boundary_intron_length,
+        current_scores, backpointers,
+    )
+
+    state = best_final_state(current_scores)
+    for position in range(record_length - 1, 0, -1):
+        path[position] = state
+        state = backpointers[position - 1, state]
+    path[0] = state
+
+
+def decode_paths_with_numba(
+    tasks,
+    transition_probabilities,
+    stop_taa=True, stop_tag=True, stop_tga=True,
+    gt_ag_intron_log_prior=GT_AG_INTRON_LOG_PRIOR,
+    gc_ag_intron_log_prior=GC_AG_INTRON_LOG_PRIOR,
+    at_ac_intron_log_prior=AT_AC_INTRON_LOG_PRIOR,
+    max_boundary_intron_length=MAX_BOUNDARY_INTRON_LENGTH_BP,
+):
+    """Decode ``[(probabilities, encoded_bases), ...]`` with the reference Numba decoder.
+
+    ``transition_probabilities`` holds one ``(bases, 4)`` array per task with
+    each transition's log-odds against staying, in nats. They are quantized to
+    the backends' int32 fixed point (``prepare_transition_probabilities``) and
+    divided by ``FIXED_POINT_SCORE_SCALE`` again, so the reference sees exactly
+    the values the native backends accumulate.
+    """
+    max_boundary_intron_length = validate_max_boundary_intron_length(max_boundary_intron_length)
+    tasks = prepare_decode_tasks(tasks)
+    transition_probabilities = prepare_transition_probabilities(transition_probabilities, tasks)
+    paths = []
+    for (probabilities, encoded_bases), fixed_point_scores in zip(tasks, transition_probabilities):
+        coding_probabilities = probabilities.astype(np.float64)
+        np.clip(coding_probabilities, MIN_PROB, MAX_PROB, out=coding_probabilities)
+        record_length = coding_probabilities.shape[0]
+
+        path = np.empty(record_length, np.int8)
+        if record_length:
+            decode_one_track(
+                np.log(coding_probabilities), np.log(1.0 - coding_probabilities),
+                fixed_point_scores.astype(np.float64) / FIXED_POINT_SCORE_SCALE,
+                np.ascontiguousarray(encoded_bases, np.int8), record_length,
+                bool(stop_taa), bool(stop_tag), bool(stop_tga),
+                float(gt_ag_intron_log_prior), float(gc_ag_intron_log_prior),
+                float(at_ac_intron_log_prior),
+                int(max_boundary_intron_length),
+                path,
+            )
+        paths.append(path)
+    return paths
+
+
+@njit(cache=True, inline="always")
+def state_emits_cds(state):
+    """Return True for states that emit CDS bases in the decoded path."""
+    return START_A <= state <= STOP
+
+
+@njit(cache=True)
+def find_gene_exon_intervals(state_path):
+    """Return exon starts, exon ends, gene ids, and gene count for one path.
+
+    This scans the path directly instead of building full-length boolean,
+    diff, and cumsum arrays. Intergenic bases between two exons start a new
+    gene; intron bases keep neighboring exons in the same gene.
+    """
+    # First pass: count how many exon intervals and genes will be emitted, so
+    # the second pass can fill fixed-size arrays without Python list growth.
+    exon_count = 0
+    gene_count = 0
+    inside_exon = False
+    seen_exon = False
+    intergenic_after_previous_exon = False
+
+    for position in range(state_path.shape[0]):
+        state = state_path[position]
+        if state_emits_cds(state):
+            if not inside_exon:
+                exon_count += 1
+                if not seen_exon:
+                    gene_count = 1
+                    seen_exon = True
+                elif intergenic_after_previous_exon:
+                    gene_count += 1
+                intergenic_after_previous_exon = False
+                inside_exon = True
+        else:
+            inside_exon = False
+            if seen_exon and state == INTERGENIC:
+                intergenic_after_previous_exon = True
+
+    exon_starts = np.empty(exon_count, np.int64)
+    exon_ends = np.empty(exon_count, np.int64)
+    gene_ids = np.empty(exon_count, np.int64)
+
+    # Second pass: record exon intervals and assign each exon to a gene.
+    exon_index = 0
+    current_gene = -1
+    inside_exon = False
+    seen_exon = False
+    intergenic_after_previous_exon = False
+
+    for position in range(state_path.shape[0]):
+        state = state_path[position]
+        if state_emits_cds(state):
+            if not inside_exon:
+                if not seen_exon:
+                    current_gene = 0
+                    seen_exon = True
+                elif intergenic_after_previous_exon:
+                    current_gene += 1
+                exon_starts[exon_index] = position
+                gene_ids[exon_index] = current_gene
+                intergenic_after_previous_exon = False
+                inside_exon = True
+        else:
+            if inside_exon:
+                exon_ends[exon_index] = position - 1
+                exon_index += 1
+                inside_exon = False
+            if seen_exon and state == INTERGENIC:
+                intergenic_after_previous_exon = True
+
+    if inside_exon:
+        exon_ends[exon_index] = state_path.shape[0] - 1
+
+    return exon_starts, exon_ends, gene_ids, gene_count
+
+
+def iter_gene_exons_with_states(state_path):
+    """Yield each gene's oriented exons and first/last CDS states."""
+    exon_starts, exon_ends, gene_ids, gene_count = find_gene_exon_intervals(state_path)
+    exon_index = 0
+    for gene_id in range(gene_count):
+        gene_exons = []
+        while exon_index < len(gene_ids) and gene_ids[exon_index] == gene_id:
+            gene_exons.append((int(exon_starts[exon_index]), int(exon_ends[exon_index])))
+            exon_index += 1
+        first_cds_state = int(state_path[gene_exons[0][0]])
+        last_cds_state = int(state_path[gene_exons[-1][1]])
+        yield gene_exons, first_cds_state, last_cds_state
+
+
+def choose_viterbi_backend(requested_backend):
+    """Return the name of the Viterbi backend to use: ``gpu``, ``cpu``, or ``numba``.
+
+    ``auto`` prefers the GPU backend, then the CPU backend, then the Numba
+    reference; ``gpu``, ``cpu``, and ``numba`` select one directly (``cpu``
+    falls back to the reference when the extension cannot be built).
+    """
+    order = {
+        "auto": ("gpu", "cpu", "numba"),
+        "gpu": ("gpu",),
+        "cpu": ("cpu", "numba"),
+        "numba": ("numba",),
+    }[requested_backend]
+
+    reasons = []
+    for backend in order:
+        if backend == "numba":
+            if requested_backend in ("auto", "cpu") and reasons:
+                warnings.warn(
+                    "CPU/GPU Viterbi backend unavailable "
+                    f"({'; '.join(reasons)}); using the Numba reference."
+                )
+            return "numba"
+
+        try:
+            module = importlib.import_module(f"viterbi_backends.decode_{backend}")
+            ok, reason = module.available()
+            if ok:
+                return backend
+            reasons.append(f"{backend}: {reason}")
+        except Exception as exc:
+            reasons.append(f"{backend}: {type(exc).__name__}: {exc}")
+
+    raise RuntimeError("no Viterbi backend available: " + "; ".join(reasons))
+
+
+def decode_paths(tasks, transition_probabilities, viterbi_backend, **options):
+    """Run the Viterbi algorithm on ``viterbi_backend`` (``gpu``, ``cpu``, or ``numba``).
+
+    Every backend takes the arguments of ``decode_paths_with_numba`` and
+    returns identical state paths.
+    """
+    if viterbi_backend == "numba":
+        return decode_paths_with_numba(tasks, transition_probabilities, **options)
+    module = importlib.import_module(f"viterbi_backends.decode_{viterbi_backend}")
+    return module.decode_paths(tasks, transition_probabilities, **options)

annotation_pipeline/src/viterbi_backends/decode_cpu.py

@@ -0,0 +1,112 @@
+"""CPU Viterbi backend: the C++/OpenMP implementation of the algorithm in viterbi.py.
+
+The C++ extension is JIT-compiled by torch.utils.cpp_extension.load on first
+use and cached under ~/.cache/torch_extensions. Tasks are
+``(P(CDS), encoded_bases)`` pairs; every task is accompanied by a
+``(bases, 4)`` array holding the transition network's log-odds against staying
+for the four permitted transitions, in nats, quantized here to the kernel's
+int32 fixed point (``prepare_transition_probabilities``). Stop
+codons, all three intron-class priors, and the boundary-intron length limit
+are explicit across every Viterbi backend.
+"""
+import os
+import sys
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from constants import (
+    AT_AC_INTRON_LOG_PRIOR,
+    GC_AG_INTRON_LOG_PRIOR,
+    GT_AG_INTRON_LOG_PRIOR,
+    MAX_BOUNDARY_INTRON_LENGTH_BP,
+)
+from utils import (
+    prepare_decode_tasks,
+    prepare_transition_probabilities,
+    validate_max_boundary_intron_length,
+    validate_reconstruction_batch_size,
+)
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_SRC = os.path.join(_HERE, "viterbi_cpu.cpp")
+_MOD = None
+
+
+def _load():
+    """Build or reuse the cached C++ extension."""
+    global _MOD
+    if _MOD is not None:
+        return _MOD
+    from torch.utils.cpp_extension import load
+
+    # Do not use -march=native: the extension must run on CPUs older than
+    # the machine that first builds it, including CPUs without AVX-512.
+    _MOD = load(
+        name="carbon_a_viterbi_cpu", sources=[_SRC],
+        extra_cflags=["-O3", "-fopenmp"],
+        extra_ldflags=["-fopenmp"],
+        verbose=False,
+    )
+    return _MOD
+
+
+def available():
+    """Return whether the CPU extension can be loaded."""
+    try:
+        _load()
+        return True, "ok"
+    except Exception as exc:
+        return False, f"{type(exc).__name__}: {exc}"
+
+
+def decode_paths(
+    tasks,
+    transition_probabilities,
+    stop_taa=True, stop_tag=True, stop_tga=True,
+    gt_ag_intron_log_prior=GT_AG_INTRON_LOG_PRIOR,
+    gc_ag_intron_log_prior=GC_AG_INTRON_LOG_PRIOR,
+    at_ac_intron_log_prior=AT_AC_INTRON_LOG_PRIOR,
+    max_boundary_intron_length=MAX_BOUNDARY_INTRON_LENGTH_BP,
+    reconstruction_batch_size=None,
+):
+    """Decode probability/base tasks with the CPU Viterbi backend.
+
+    ``transition_probabilities`` holds one ``(bases, 4)`` array per task of
+    nat-valued log-odds in ``TRANSITION_NAMES`` column order.
+    ``reconstruction_batch_size`` is the number of bases processed together
+    during final path reconstruction. ``None`` uses the backend default.
+    """
+    max_boundary_intron_length = validate_max_boundary_intron_length(max_boundary_intron_length)
+    reconstruction_batch_size = validate_reconstruction_batch_size(reconstruction_batch_size)
+    tasks = prepare_decode_tasks(tasks)
+    transition_probabilities = prepare_transition_probabilities(transition_probabilities, tasks)
+    if not tasks:
+        return []
+    module = _load()
+    record_offsets = np.zeros(len(tasks) + 1, np.int64)
+    for index, (_, encoded_bases) in enumerate(tasks):
+        record_offsets[index + 1] = record_offsets[index] + len(encoded_bases)
+
+    probability_bits = np.ascontiguousarray(
+        np.concatenate([np.asarray(probabilities, np.float16) for probabilities, _ in tasks])
+    ).view(np.uint16)
+    encoded_bases = np.ascontiguousarray(
+        np.concatenate([np.asarray(bases, np.int8) for _, bases in tasks]),
+        dtype=np.int8,
+    )
+    packed_transition_probabilities = np.ascontiguousarray(np.concatenate(transition_probabilities), dtype=np.int32)
+    combined_path = np.asarray(
+        module.decode(
+            probability_bits, encoded_bases, packed_transition_probabilities, record_offsets,
+            int(stop_taa), int(stop_tag), int(stop_tga),
+            float(gt_ag_intron_log_prior), float(gc_ag_intron_log_prior),
+            float(at_ac_intron_log_prior),
+            int(max_boundary_intron_length),
+            0 if reconstruction_batch_size is None else reconstruction_batch_size,
+        )
+    )
+    return [
+        combined_path[int(record_offsets[index]):int(record_offsets[index + 1])]
+        for index in range(len(tasks))
+    ]

annotation_pipeline/src/viterbi_backends/decode_gpu.py

@@ -0,0 +1,167 @@
+"""GPU Viterbi backend: the CUDA implementation of the algorithm in viterbi.py.
+
+The CUDA extension is JIT-compiled by torch.utils.cpp_extension.load on first
+use. Tasks are ``(P(CDS), encoded_bases)`` pairs; every task is accompanied by
+a ``(bases, 4)`` array holding the transition network's log-odds against
+staying for the four permitted transitions, in nats, quantized here to the
+kernel's int32 fixed point (``prepare_transition_probabilities``). Stop
+codons, all three intron-class priors, and the boundary-intron length limit
+are explicit across every Viterbi backend.
+
+decode_paths packs tasks into sub-batches sized by available VRAM. The kernel
+windows the backpointer array internally, so a single long chromosome can be
+decoded without materializing a genome-sized backpointer buffer.
+"""
+import os
+import sys
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from constants import (
+    AT_AC_INTRON_LOG_PRIOR,
+    GC_AG_INTRON_LOG_PRIOR,
+    GPU_ESTIMATED_BYTES_PER_BASE,
+    GPU_MAX_BATCH_BASES,
+    GPU_MEMORY_FRACTION,
+    GPU_MIN_BATCH_BASES,
+    GPU_SEGMENT_LENGTH,
+    GT_AG_INTRON_LOG_PRIOR,
+    MAX_BOUNDARY_INTRON_LENGTH_BP,
+)
+from utils import (
+    prepare_decode_tasks,
+    prepare_transition_probabilities,
+    validate_max_boundary_intron_length,
+    validate_reconstruction_batch_size,
+)
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_SRC = os.path.join(_HERE, "viterbi_gpu.cu")
+_MOD = None
+
+
+def _load():
+    """Build or reuse the cached CUDA extension."""
+    global _MOD
+    if _MOD is not None:
+        return _MOD
+    from torch.utils.cpp_extension import load
+
+    _MOD = load(
+        name="carbon_a_viterbi_gpu", sources=[_SRC],
+        with_cuda=True, extra_cuda_cflags=["-O3"], verbose=False,
+    )
+    return _MOD
+
+
+def available():
+    """Return whether CUDA and the GPU extension are available."""
+    try:
+        import torch
+        if not torch.cuda.is_available():
+            return False, "no CUDA device"
+        _load()
+        return True, "ok"
+    except Exception as exc:
+        return False, f"{type(exc).__name__}: {exc}"
+
+
+def batch_base_budget(free_bytes):
+    """Return a VRAM-based sub-batch size that also fits CUDA int32 indexing."""
+    memory_budget = max(
+        GPU_MIN_BATCH_BASES,
+        int(free_bytes * GPU_MEMORY_FRACTION / GPU_ESTIMATED_BYTES_PER_BASE),
+    )
+    return min(memory_budget, GPU_MAX_BATCH_BASES)
+
+
+def _decode_batch(module, tasks, transition_probabilities, reconstruction_batch_size):
+    """Pack and decode one GPU sub-batch."""
+    record_offsets = np.zeros(len(tasks) + 1, np.int64)
+    for index, (_, encoded_bases) in enumerate(tasks):
+        record_offsets[index + 1] = record_offsets[index] + len(encoded_bases)
+
+    probability_bits = np.ascontiguousarray(
+        np.concatenate([np.asarray(probabilities, np.float16) for probabilities, _ in tasks])
+    ).view(np.uint16)
+    encoded_bases = np.ascontiguousarray(
+        np.concatenate([np.asarray(bases, np.int8) for _, bases in tasks]),
+        dtype=np.int8,
+    )
+    packed_transition_probabilities = np.ascontiguousarray(np.concatenate(transition_probabilities), dtype=np.int32)
+    combined_path = np.asarray(
+        module.decode(
+            probability_bits, encoded_bases, packed_transition_probabilities, record_offsets,
+            GPU_SEGMENT_LENGTH,
+            0 if reconstruction_batch_size is None else reconstruction_batch_size,
+        )
+    )
+    return [
+        combined_path[int(record_offsets[index]):int(record_offsets[index + 1])]
+        for index in range(len(tasks))
+    ]
+
+
+def decode_paths(
+    tasks,
+    transition_probabilities,
+    stop_taa=True, stop_tag=True, stop_tga=True,
+    gt_ag_intron_log_prior=GT_AG_INTRON_LOG_PRIOR,
+    gc_ag_intron_log_prior=GC_AG_INTRON_LOG_PRIOR,
+    at_ac_intron_log_prior=AT_AC_INTRON_LOG_PRIOR,
+    max_boundary_intron_length=MAX_BOUNDARY_INTRON_LENGTH_BP,
+    reconstruction_batch_size=None,
+):
+    """Decode probability/base tasks with the GPU Viterbi backend.
+
+    ``transition_probabilities`` holds one ``(bases, 4)`` array per task of
+    nat-valued log-odds in ``TRANSITION_NAMES`` column order.
+    ``reconstruction_batch_size`` is the number of bases processed together
+    during final path reconstruction. ``None`` lets CUDA size it from VRAM.
+    """
+    max_boundary_intron_length = validate_max_boundary_intron_length(max_boundary_intron_length)
+    reconstruction_batch_size = validate_reconstruction_batch_size(reconstruction_batch_size)
+    tasks = prepare_decode_tasks(tasks)
+    transition_probabilities = prepare_transition_probabilities(transition_probabilities, tasks)
+    if not tasks:
+        return []
+    for _, encoded_bases in tasks:
+        if len(encoded_bases) > GPU_MAX_BATCH_BASES:
+            raise ValueError(
+                f"one GPU Viterbi backend record has {len(encoded_bases)} bases; "
+                f"the maximum supported length is {GPU_MAX_BATCH_BASES}"
+            )
+    import torch
+
+    module = _load()
+    module.init(
+        int(stop_taa), int(stop_tag), int(stop_tga),
+        float(gt_ag_intron_log_prior), float(gc_ag_intron_log_prior),
+        float(at_ac_intron_log_prior),
+        int(max_boundary_intron_length),
+    )
+
+    free_bytes = torch.cuda.mem_get_info()[0]
+    base_budget = batch_base_budget(free_bytes)
+
+    decoded_paths = []
+    batch_start = 0
+    while batch_start < len(tasks):
+        batch_end = batch_start
+        batch_bases = 0
+        while batch_end < len(tasks):
+            next_record_bases = len(tasks[batch_end][1])
+            if batch_end > batch_start and batch_bases + next_record_bases > base_budget:
+                break
+            batch_bases += next_record_bases
+            batch_end += 1
+
+        decoded_paths.extend(
+            _decode_batch(
+                module, tasks[batch_start:batch_end],
+                transition_probabilities[batch_start:batch_end], reconstruction_batch_size,
+            )
+        )
+        batch_start = batch_end
+    return decoded_paths

annotation_pipeline/src/viterbi_backends/viterbi_cpu.cpp

@@ -0,0 +1,806 @@
+// viterbi_cpu.cpp -- CPU Viterbi backend: C++/OpenMP implementation of the algorithm in viterbi.py.
+//
+// This file implements the same 30-state decoder grammar as viterbi.py:
+//   - six exon prefix types for stop-codon detection;
+//   - GT/GC-AG introns enter FIRST, then move to READY before AG exit;
+//   - AT-AC introns use one state table because AT cannot overlap AC;
+//   - Python supplies explicit GT-AG, GC-AG, and AT-AC log priors;
+//   - N freely matches any concrete base in codons and splice signals;
+//   - dedicated leading/trailing states represent introns crossing a record edge;
+//   - the transition network supplies the probabilities of the four permitted
+//     transitions per base (enter gene, leave gene, enter intron, leave
+//     intron) as log-odds in the decoder's fixed-point units;
+//   - stop read-through is disabled.
+//
+// Per-base log-probability and transition increments are int32 fixed-point
+// values from a fp16 log LUT. Accumulated path scores use int64 and are kept
+// unnormalized, so this implementation does not discard a state just because
+// it temporarily falls behind the current best state. Long records keep score
+// checkpoints at reconstruction-batch boundaries, so backpointer memory does
+// not grow with the complete record length.
+#include <cstdint>
+#include <cmath>
+#include <cstring>
+#include <stdexcept>
+#include <vector>
+#include <omp.h>
+
+using AccumulatedScore = int64_t;
+using StepScore = int32_t;
+static const double SCORE_SCALE = 16384.0;
+static inline StepScore score_from_nats(double nats) {
+    return static_cast<StepScore>(std::llround(nats * SCORE_SCALE));
+}
+// Fixed-point code uses this sentinel in place of IEEE -inf.
+static const AccumulatedScore UNREACHABLE_ACCUMULATED_SCORE = -4'000'000'000'000'000'000LL;
+static const int64_t DEFAULT_RECONSTRUCTION_BATCH_SIZE = 8'000'000;
+
+static const int STATE_COUNT = 30;
+static const int PREFIX_TYPE_COUNT = 6;
+static const int INTERGENIC = 0;
+static const int START_A = 1;
+static const int START_AT = 2;
+static const int EXON_PREFIX_TYPE_0 = 3;
+static const int STOP = 9;
+static const int GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0 = 10;
+static const int GT_GC_AG_INTRON_READY_PREFIX_TYPE_0 = 16;
+static const int AT_AC_INTRON_PREFIX_TYPE_0 = 22;
+static const int LEADING_BOUNDARY_INTRON = 28;
+static const int TRAILING_BOUNDARY_INTRON = 29;
+
+// State layout matches viterbi.py and viterbi_gpu.cu:
+//   0 intergenic; 1..2 ATG start; 3..8 exon prefix types; 9 stop;
+//   10..15 GT/GC-AG FIRST; 16..21 GT/GC-AG READY; 22..27 AT-AC;
+//   28 leading boundary intron; 29 trailing boundary intron.
+
+static const int BASE_A = 0;
+static const int BASE_T = 1;
+static const int BASE_C = 2;
+static const int BASE_G = 3;
+static const int BASE_N = 4;
+
+// Transition probability tracks per base, in the column order of constants.py.
+// They arrive from Python as log-odds against staying, already quantized to
+// the decoder's fixed-point units, so they add directly to path scores.
+static const int ENTER_GENE = 0;
+static const int LEAVE_GENE = 1;
+static const int ENTER_INTRON = 2;
+static const int LEAVE_INTRON = 3;
+static const int TRANSITION_COUNT = 4;
+
+static inline bool is_exon_state(int state) {
+    return state >= EXON_PREFIX_TYPE_0
+        && state < EXON_PREFIX_TYPE_0 + PREFIX_TYPE_COUNT;
+}
+
+static inline bool is_gt_gc_intron_entry_state(int state) {
+    return state >= GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0
+        && state < GT_GC_AG_INTRON_READY_PREFIX_TYPE_0;
+}
+
+static inline bool is_gt_gc_acceptor_state(int state) {
+    return state >= GT_GC_AG_INTRON_READY_PREFIX_TYPE_0
+        && state < AT_AC_INTRON_PREFIX_TYPE_0;
+}
+
+static inline bool is_at_ac_intron_state(int state) {
+    return state >= AT_AC_INTRON_PREFIX_TYPE_0
+        && state < LEADING_BOUNDARY_INTRON;
+}
+
+// Transition log-odds of one legal state change. Each value is attached to the
+// base where the new segment begins: the first start-codon base (enter gene),
+// the last stop-codon base (leave gene, charged on the base that closes the
+// gene), the first intron base (enter intron), or the first exon base after an
+// intron (leave intron). Introns that cross a record edge are entered and left
+// the same way.
+static inline StepScore transition_log_odds(
+    int source_state, int target_state, const StepScore* position_transition_probabilities
+) {
+    if (source_state == INTERGENIC && target_state == START_A) {
+        return position_transition_probabilities[ENTER_GENE];
+    }
+    // Leaving the gene takes precedence when an intron exit completes a stop
+    // codon that was split by the intron.
+    if (target_state == STOP) return position_transition_probabilities[LEAVE_GENE];
+    if (is_exon_state(source_state)
+        && (is_gt_gc_intron_entry_state(target_state)
+            || is_at_ac_intron_state(target_state)
+            || target_state == TRAILING_BOUNDARY_INTRON)) {
+        return position_transition_probabilities[ENTER_INTRON];
+    }
+    if ((is_gt_gc_acceptor_state(source_state)
+            || is_at_ac_intron_state(source_state)
+            || source_state == LEADING_BOUNDARY_INTRON)
+        && is_exon_state(target_state)) {
+        return position_transition_probabilities[LEAVE_INTRON];
+    }
+    return StepScore(0);
+}
+
+static inline bool base_matches(int observed_base, int expected_base) {
+    return observed_base == expected_base || observed_base == BASE_N;
+}
+
+static inline bool is_exact_enabled_stop_codon(
+    int first_base, int second_base, int third_base,
+    bool stop_taa, bool stop_tag, bool stop_tga
+) {
+    if (first_base != BASE_T) return false;
+    if (second_base == BASE_A) {
+        if (third_base == BASE_A && stop_taa) return true;
+        if (third_base == BASE_G && stop_tag) return true;
+    }
+    return second_base == BASE_G
+        && third_base == BASE_A
+        && stop_tga;
+}
+
+static StepScore coding_score_lookup[65536], noncoding_score_lookup[65536];
+
+static double half_bits_to_double(uint16_t half_bits) {
+    uint32_t sign = static_cast<uint32_t>(half_bits & 0x8000) << 16;
+    uint32_t exponent = (half_bits >> 10) & 0x1F;
+    uint32_t mantissa = half_bits & 0x3FF;
+    uint32_t float_bits;
+    if (exponent == 0) {
+        if (mantissa == 0) {
+            float_bits = sign;
+        } else {
+            int normalization_shift = -1;
+            do {
+                normalization_shift++;
+                mantissa <<= 1;
+            } while (!(mantissa & 0x400));
+            float_bits = sign
+                | (static_cast<uint32_t>(127 - 15 - normalization_shift) << 23)
+                | ((mantissa & 0x3FF) << 13);
+        }
+    } else if (exponent == 0x1F) {
+        float_bits = sign | 0x7F800000u | (mantissa << 13);
+    } else {
+        float_bits = sign | ((exponent + (127 - 15)) << 23) | (mantissa << 13);
+    }
+    float value;
+    memcpy(&value, &float_bits, sizeof(value));
+    return static_cast<double>(value);
+}
+
+static bool build_log_score_lookup() {
+    for (int index = 0; index < 65536; index++) {
+        double probability = half_bits_to_double(static_cast<uint16_t>(index));
+        if (probability < 1e-6) probability = 1e-6;
+        else if (probability > 1.0 - 1e-6) probability = 1.0 - 1e-6;
+        coding_score_lookup[index] = score_from_nats(log(probability));
+        noncoding_score_lookup[index] = score_from_nats(log(1.0 - probability));
+    }
+    return true;
+}
+
+static inline int exon_state_for_prefix(int prefix_type) {
+    return EXON_PREFIX_TYPE_0 + prefix_type;
+}
+
+static inline int next_coding_state_after_base(
+    int prefix_type, int base,
+    bool stop_taa, bool stop_tag, bool stop_tga
+) {
+    if (prefix_type == 0) return base == BASE_T ? exon_state_for_prefix(1) : exon_state_for_prefix(2);
+    if (prefix_type == 1) {
+        if (base == BASE_A) return exon_state_for_prefix(3);
+        if (base == BASE_G) return exon_state_for_prefix(4);
+        return exon_state_for_prefix(5);
+    }
+    if (prefix_type == 2) return exon_state_for_prefix(5);
+    if (prefix_type == 3) {
+        if (base == BASE_A && stop_taa) return STOP;
+        if (base == BASE_G && stop_tag) return STOP;
+        return exon_state_for_prefix(0);
+    }
+    if (prefix_type == 4) {
+        if (base == BASE_A && stop_tga) return STOP;
+        return exon_state_for_prefix(0);
+    }
+    return exon_state_for_prefix(0);
+}
+
+static inline void keep_best_transition(
+    AccumulatedScore* next_scores, int8_t* predecessors,
+    int target_state, AccumulatedScore candidate_score, int source_state,
+    const StepScore* position_transition_probabilities
+) {
+    candidate_score += transition_log_odds(
+        source_state, target_state, position_transition_probabilities
+    );
+    if (candidate_score > next_scores[target_state]) {
+        next_scores[target_state] = candidate_score;
+        predecessors[target_state] = static_cast<int8_t>(source_state);
+    }
+}
+
+static inline void emit_coding_base(
+    const AccumulatedScore* current_scores, AccumulatedScore* next_scores,
+    int8_t* predecessors, int source_state, int prefix_type,
+    StepScore coding_score, const StepScore* position_transition_probabilities, int base,
+    bool stop_taa, bool stop_tag, bool stop_tga
+) {
+    AccumulatedScore source_score = current_scores[source_state];
+    if (source_score == UNREACHABLE_ACCUMULATED_SCORE) return;
+    AccumulatedScore candidate_score = source_score + coding_score;
+    if (base == BASE_N) {
+        for (int resolved_base = BASE_A; resolved_base <= BASE_G; resolved_base++) {
+            int target_state = next_coding_state_after_base(
+                prefix_type, resolved_base, stop_taa, stop_tag, stop_tga
+            );
+            keep_best_transition(
+                next_scores, predecessors, target_state, candidate_score, source_state,
+                position_transition_probabilities
+            );
+        }
+    } else {
+        int target_state = next_coding_state_after_base(
+            prefix_type, base, stop_taa, stop_tag, stop_tga
+        );
+        keep_best_transition(
+            next_scores, predecessors, target_state, candidate_score, source_state,
+            position_transition_probabilities
+        );
+    }
+}
+
+// Transitions across the left record edge are charged like interior ones: a
+// start codon beginning at the first base is charged enter_gene[0] and a stop
+// codon completed there is charged leave_gene[0]; every other first-base state
+// is charged nothing.
+static void initialize_first_base_scores(
+    const uint16_t* probability_bits, const StepScore* transition_probabilities,
+    const int8_t* bases, int record_length,
+    bool stop_taa, bool stop_tag, bool stop_tga,
+    bool boundary_introns_enabled, AccumulatedScore* current_scores
+) {
+    StepScore coding_score = coding_score_lookup[probability_bits[0]];
+    StepScore noncoding_score = noncoding_score_lookup[probability_bits[0]];
+    int base = bases[0];
+
+    for (int state = 0; state < STATE_COUNT; state++) {
+        current_scores[state] = UNREACHABLE_ACCUMULATED_SCORE;
+    }
+
+    current_scores[INTERGENIC] = noncoding_score;
+    current_scores[exon_state_for_prefix(0)] = coding_score;
+    current_scores[exon_state_for_prefix(5)] = coding_score;
+    if (base == BASE_N) {
+        for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+            current_scores[exon_state_for_prefix(prefix_type)] = coding_score;
+        }
+    } else {
+        int one_base_prefix_type = base == BASE_T ? 1 : 2;
+        current_scores[exon_state_for_prefix(one_base_prefix_type)] = coding_score;
+        if (base == BASE_A) current_scores[exon_state_for_prefix(3)] = coding_score;
+        else if (base == BASE_G) current_scores[exon_state_for_prefix(4)] = coding_score;
+    }
+    if ((base_matches(base, BASE_A) && (stop_taa || stop_tga))
+        || (base_matches(base, BASE_G) && stop_tag)) {
+        current_scores[STOP] = static_cast<AccumulatedScore>(coding_score)
+            + transition_probabilities[LEAVE_GENE];
+    }
+
+    if (boundary_introns_enabled) {
+        current_scores[LEADING_BOUNDARY_INTRON] = noncoding_score;
+    }
+
+    bool second_start_base_matches = record_length == 1 || base_matches(bases[1], BASE_T);
+    bool third_start_base_matches = record_length <= 2 || base_matches(bases[2], BASE_G);
+    bool has_atg_start = base_matches(base, BASE_A)
+        && second_start_base_matches
+        && third_start_base_matches;
+    if (has_atg_start) {
+        current_scores[START_A] = static_cast<AccumulatedScore>(coding_score)
+            + transition_probabilities[ENTER_GENE];
+    }
+}
+
+static void run_viterbi_forward(
+    const uint16_t* probability_bits, const int8_t* bases,
+    const StepScore* transition_probabilities, int record_length,
+    int segment_start, int segment_end, bool is_record_start,
+    StepScore gt_ag_intron_score, StepScore gc_ag_intron_score, StepScore at_ac_intron_score,
+    int max_boundary_intron_length,
+    bool stop_taa, bool stop_tag, bool stop_tga,
+    int8_t* backpointers, AccumulatedScore* output_scores,
+    const AccumulatedScore* input_scores = nullptr
+) {
+    AccumulatedScore current_buffer[STATE_COUNT];
+    AccumulatedScore next_buffer[STATE_COUNT];
+    AccumulatedScore* current_scores = current_buffer;
+    AccumulatedScore* next_scores = next_buffer;
+
+    if (is_record_start) {
+        initialize_first_base_scores(
+            probability_bits, transition_probabilities, bases, record_length,
+            stop_taa, stop_tag, stop_tga,
+            max_boundary_intron_length != 0,
+            current_scores
+        );
+    } else {
+        for (int state = 0; state < STATE_COUNT; state++) {
+            current_scores[state] = input_scores[state];
+        }
+    }
+
+    for (int position = segment_start + 1; position <= segment_end; position++) {
+        int base = bases[position];
+        StepScore coding_score = coding_score_lookup[probability_bits[position]];
+        StepScore noncoding_score = noncoding_score_lookup[probability_bits[position]];
+        const StepScore* position_transition_probabilities =
+            transition_probabilities + static_cast<size_t>(position) * TRANSITION_COUNT;
+
+        // Missing motif bases at record edges are unknown. Internal positions
+        // still require the complete start, donor, or acceptor motif.
+        bool next_base_is_missing = position + 1 == record_length;
+        int next_base = next_base_is_missing ? -1 : bases[position + 1];
+        bool has_gt_ag_donor = base_matches(base, BASE_G)
+            && (next_base_is_missing || base_matches(next_base, BASE_T));
+        bool has_gc_ag_donor = base_matches(base, BASE_G)
+            && (next_base_is_missing || base_matches(next_base, BASE_C));
+        bool acceptor_first_base_matches = position == 1
+            || base_matches(bases[position - 2], BASE_A);
+        bool has_ag_acceptor = acceptor_first_base_matches
+            && base_matches(bases[position - 1], BASE_G);
+        bool has_ac_acceptor = acceptor_first_base_matches
+            && base_matches(bases[position - 1], BASE_C);
+        bool third_start_base_matches = position + 2 >= record_length
+            || base_matches(bases[position + 2], BASE_G);
+        bool has_atg_start = base_matches(base, BASE_A)
+            && (next_base_is_missing || base_matches(next_base, BASE_T))
+            && third_start_base_matches;
+        // Boundary guards require exact five-base patterns. N and missing
+        // edge context remain valid for ordinary motif matching only.
+        bool has_exact_acceptor_followed_by_atg = position >= 2
+            && position + 2 < record_length
+            && bases[position - 2] == BASE_A
+            && (bases[position - 1] == BASE_G || bases[position - 1] == BASE_C)
+            && base == BASE_A
+            && next_base == BASE_T
+            && bases[position + 2] == BASE_G;
+        bool has_at_ac_donor = base_matches(base, BASE_A)
+            && (next_base_is_missing || base_matches(next_base, BASE_T));
+        bool has_any_donor = has_gt_ag_donor || has_gc_ag_donor || has_at_ac_donor;
+        bool has_exact_donor = position + 1 < record_length
+            && ((base == BASE_G && (next_base == BASE_T || next_base == BASE_C))
+                || (base == BASE_A && next_base == BASE_T));
+        bool has_exact_enabled_stop_followed_by_donor = false;
+        if (has_exact_donor && position >= 3) {
+            has_exact_enabled_stop_followed_by_donor = is_exact_enabled_stop_codon(
+                bases[position - 3], bases[position - 2], bases[position - 1],
+                stop_taa, stop_tag, stop_tga
+            );
+        }
+        // A leading self-loop into position makes its run position + 1 bases.
+        // A trailing entry here emits record_length - position bases through
+        // the right edge, so both limits follow directly from coordinates.
+        bool boundary_introns_enabled = max_boundary_intron_length != 0;
+        bool boundary_intron_length_is_unlimited = max_boundary_intron_length == -1;
+        bool can_extend_leading_boundary_intron =
+            boundary_introns_enabled
+            && (boundary_intron_length_is_unlimited
+                || position < max_boundary_intron_length);
+        bool can_enter_trailing_boundary_intron =
+            boundary_introns_enabled
+            && (boundary_intron_length_is_unlimited
+                || record_length - position <= max_boundary_intron_length);
+
+        int8_t local_predecessors[STATE_COUNT];
+        int8_t* predecessors = backpointers
+            ? &backpointers[
+                static_cast<size_t>(position - segment_start - 1) * STATE_COUNT
+            ]
+            : local_predecessors;
+        for (int state = 0; state < STATE_COUNT; state++) {
+            next_scores[state] = UNREACHABLE_ACCUMULATED_SCORE;
+            predecessors[state] = 0;
+        }
+
+        // Intergenic, stop, and start-codon transitions.
+        if (current_scores[INTERGENIC] != UNREACHABLE_ACCUMULATED_SCORE) {
+            keep_best_transition(
+                next_scores, predecessors, INTERGENIC,
+                current_scores[INTERGENIC] + noncoding_score, INTERGENIC,
+                position_transition_probabilities
+            );
+            if (has_atg_start) {
+                keep_best_transition(
+                    next_scores, predecessors, START_A,
+                    current_scores[INTERGENIC] + coding_score,
+                    INTERGENIC,
+                    position_transition_probabilities
+                );
+            }
+        }
+
+        if (current_scores[STOP] != UNREACHABLE_ACCUMULATED_SCORE) {
+            keep_best_transition(
+                next_scores, predecessors, INTERGENIC,
+                current_scores[STOP] + noncoding_score, STOP,
+                position_transition_probabilities
+            );
+        }
+
+        if (current_scores[START_A] != UNREACHABLE_ACCUMULATED_SCORE) {
+            keep_best_transition(
+                next_scores, predecessors, START_AT,
+                current_scores[START_A] + coding_score, START_A,
+                position_transition_probabilities
+            );
+        }
+        if (current_scores[START_AT] != UNREACHABLE_ACCUMULATED_SCORE) {
+            keep_best_transition(
+                next_scores, predecessors, exon_state_for_prefix(0),
+                current_scores[START_AT] + coding_score, START_AT,
+                position_transition_probabilities
+            );
+        }
+
+        // Every prefix type is allowed when the leading boundary intron exits,
+        // so the state carries none. Only an exact acceptor followed by an
+        // exact ATG activates the guard; N leaves both paths available.
+        if (current_scores[LEADING_BOUNDARY_INTRON] != UNREACHABLE_ACCUMULATED_SCORE) {
+            if (can_extend_leading_boundary_intron) {
+                keep_best_transition(
+                    next_scores, predecessors, LEADING_BOUNDARY_INTRON,
+                    current_scores[LEADING_BOUNDARY_INTRON] + noncoding_score,
+                    LEADING_BOUNDARY_INTRON,
+                    position_transition_probabilities
+                );
+            }
+            bool can_exit_leading_intron =
+                (has_ag_acceptor || has_ac_acceptor)
+                && !has_exact_acceptor_followed_by_atg;
+            if (can_exit_leading_intron) {
+                for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+                    emit_coding_base(
+                        current_scores, next_scores, predecessors,
+                        LEADING_BOUNDARY_INTRON, prefix_type,
+                        coding_score, position_transition_probabilities,
+                        base, stop_taa, stop_tag, stop_tga
+                    );
+                }
+            }
+        }
+
+        // Exons emit coding bases, open an ordinary internal intron, or enter
+        // the trailing state when the intron continues beyond the record.
+        for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+            int exon_state = exon_state_for_prefix(prefix_type);
+            AccumulatedScore exon_score = current_scores[exon_state];
+            if (exon_score == UNREACHABLE_ACCUMULATED_SCORE) continue;
+
+            emit_coding_base(
+                current_scores, next_scores, predecessors,
+                exon_state, prefix_type,
+                coding_score, position_transition_probabilities,
+                base, stop_taa, stop_tag, stop_tga
+            );
+
+            if (has_gt_ag_donor) {
+                keep_best_transition(
+                    next_scores, predecessors,
+                    GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0 + prefix_type,
+                    exon_score + noncoding_score + gt_ag_intron_score,
+                    exon_state,
+                    position_transition_probabilities
+                );
+            }
+            if (has_gc_ag_donor) {
+                keep_best_transition(
+                    next_scores, predecessors,
+                    GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0 + prefix_type,
+                    exon_score + noncoding_score + gc_ag_intron_score,
+                    exon_state,
+                    position_transition_probabilities
+                );
+            }
+            if (has_at_ac_donor) {
+                keep_best_transition(
+                    next_scores, predecessors, AT_AC_INTRON_PREFIX_TYPE_0 + prefix_type,
+                    exon_score + noncoding_score + at_ac_intron_score,
+                    exon_state,
+                    position_transition_probabilities
+                );
+            }
+            if (has_any_donor
+                && !has_exact_enabled_stop_followed_by_donor
+                && can_enter_trailing_boundary_intron) {
+                if (has_gt_ag_donor) {
+                    keep_best_transition(
+                        next_scores, predecessors, TRAILING_BOUNDARY_INTRON,
+                        exon_score + noncoding_score + gt_ag_intron_score,
+                        exon_state,
+                        position_transition_probabilities
+                    );
+                }
+                if (has_gc_ag_donor) {
+                    keep_best_transition(
+                        next_scores, predecessors, TRAILING_BOUNDARY_INTRON,
+                        exon_score + noncoding_score + gc_ag_intron_score,
+                        exon_state,
+                        position_transition_probabilities
+                    );
+                }
+                if (has_at_ac_donor) {
+                    keep_best_transition(
+                        next_scores, predecessors, TRAILING_BOUNDARY_INTRON,
+                        exon_score + noncoding_score + at_ac_intron_score,
+                        exon_state,
+                        position_transition_probabilities
+                    );
+                }
+            }
+        }
+
+        // FIRST consumes one additional base before AG exit is legal.
+        for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+            int intron_state = GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0 + prefix_type;
+            if (current_scores[intron_state] == UNREACHABLE_ACCUMULATED_SCORE) continue;
+            keep_best_transition(
+                next_scores, predecessors,
+                GT_GC_AG_INTRON_READY_PREFIX_TYPE_0 + prefix_type,
+                current_scores[intron_state] + noncoding_score, intron_state,
+                position_transition_probabilities
+            );
+        }
+
+        // READY self-loops or exits a GT/GC-AG intron after AG.
+        for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+            int intron_state = GT_GC_AG_INTRON_READY_PREFIX_TYPE_0 + prefix_type;
+            if (current_scores[intron_state] == UNREACHABLE_ACCUMULATED_SCORE) continue;
+            keep_best_transition(
+                next_scores, predecessors, intron_state,
+                current_scores[intron_state] + noncoding_score, intron_state,
+                position_transition_probabilities
+            );
+            if (has_ag_acceptor) {
+                emit_coding_base(
+                    current_scores, next_scores, predecessors,
+                    intron_state, prefix_type,
+                    coding_score, position_transition_probabilities,
+                    base, stop_taa, stop_tag, stop_tga
+                );
+            }
+        }
+
+        // AT-AC introns self-loop or exit after AC.
+        for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+            int intron_state = AT_AC_INTRON_PREFIX_TYPE_0 + prefix_type;
+            if (current_scores[intron_state] == UNREACHABLE_ACCUMULATED_SCORE) continue;
+            keep_best_transition(
+                next_scores, predecessors, intron_state,
+                current_scores[intron_state] + noncoding_score, intron_state,
+                position_transition_probabilities
+            );
+            if (has_ac_acceptor) {
+                emit_coding_base(
+                    current_scores, next_scores, predecessors,
+                    intron_state, prefix_type,
+                    coding_score, position_transition_probabilities,
+                    base, stop_taa, stop_tag, stop_tga
+                );
+            }
+        }
+
+        // A trailing boundary intron cannot return to an exon.
+        if (current_scores[TRAILING_BOUNDARY_INTRON] != UNREACHABLE_ACCUMULATED_SCORE) {
+            keep_best_transition(
+                next_scores, predecessors, TRAILING_BOUNDARY_INTRON,
+                current_scores[TRAILING_BOUNDARY_INTRON] + noncoding_score,
+                TRAILING_BOUNDARY_INTRON,
+                position_transition_probabilities
+            );
+        }
+
+        AccumulatedScore* previous_scores = current_scores;
+        current_scores = next_scores;
+        next_scores = previous_scores;
+    }
+
+    if (output_scores) {
+        for (int state = 0; state < STATE_COUNT; state++) {
+            output_scores[state] = current_scores[state];
+        }
+    }
+}
+
+static int best_final_state(const AccumulatedScore* current_scores) {
+    int best_state = INTERGENIC;
+    for (int state = START_A; state <= STOP; state++) {
+        if (current_scores[state] > current_scores[best_state]) best_state = state;
+    }
+    if (current_scores[TRAILING_BOUNDARY_INTRON] > current_scores[best_state]) {
+        best_state = TRAILING_BOUNDARY_INTRON;
+    }
+    return best_state;
+}
+
+static void decode_one_track(
+    const uint16_t* probability_bits, const int8_t* bases,
+    const StepScore* transition_probabilities, int record_length,
+    StepScore gt_ag_intron_score, StepScore gc_ag_intron_score, StepScore at_ac_intron_score,
+    int max_boundary_intron_length,
+    bool stop_taa, bool stop_tag, bool stop_tga,
+    int64_t reconstruction_batch_size, int8_t* path,
+    std::vector<int8_t>& backpointer_buffer,
+    std::vector<AccumulatedScore>& checkpoints
+) {
+    if (record_length <= 0) return;
+
+    if (static_cast<int64_t>(record_length) <= reconstruction_batch_size) {
+        backpointer_buffer.assign(static_cast<size_t>(record_length) * STATE_COUNT, 0);
+        AccumulatedScore final_scores[STATE_COUNT];
+        run_viterbi_forward(
+            probability_bits, bases, transition_probabilities, record_length, 0, record_length - 1, true,
+            gt_ag_intron_score, gc_ag_intron_score, at_ac_intron_score,
+            max_boundary_intron_length,
+            stop_taa, stop_tag, stop_tga,
+            backpointer_buffer.data(), final_scores
+        );
+
+        int state = best_final_state(final_scores);
+        for (int64_t position = record_length - 1; position >= 1; position--) {
+            path[position] = static_cast<int8_t>(state);
+            state = static_cast<int>(
+                backpointer_buffer[static_cast<size_t>(position - 1) * STATE_COUNT + state]
+            );
+        }
+        path[0] = static_cast<int8_t>(state);
+        return;
+    }
+
+    int64_t window_count = (
+        record_length - 1 + reconstruction_batch_size - 1
+    ) / reconstruction_batch_size;
+    if (window_count < 1) window_count = 1;
+    checkpoints.assign(static_cast<size_t>(window_count + 1) * STATE_COUNT, 0);
+
+    // First pass keeps only the score vector at each window boundary.
+    AccumulatedScore current_scores[STATE_COUNT];
+    initialize_first_base_scores(
+        probability_bits, transition_probabilities, bases, record_length,
+        stop_taa, stop_tag, stop_tga,
+        max_boundary_intron_length != 0,
+        current_scores
+    );
+    for (int64_t window = 0; window < window_count; window++) {
+        for (int state = 0; state < STATE_COUNT; state++) {
+            checkpoints[static_cast<size_t>(window) * STATE_COUNT + state] =
+                current_scores[state];
+        }
+        int64_t segment_start = window * reconstruction_batch_size;
+        int64_t segment_end = segment_start + reconstruction_batch_size;
+        if (segment_end > record_length - 1) segment_end = record_length - 1;
+        run_viterbi_forward(
+            probability_bits, bases, transition_probabilities, record_length,
+            static_cast<int>(segment_start), static_cast<int>(segment_end), window == 0,
+            gt_ag_intron_score, gc_ag_intron_score, at_ac_intron_score,
+            max_boundary_intron_length,
+            stop_taa, stop_tag, stop_tga,
+            nullptr, current_scores, current_scores
+        );
+    }
+    for (int state = 0; state < STATE_COUNT; state++) {
+        checkpoints[static_cast<size_t>(window_count) * STATE_COUNT + state] =
+            current_scores[state];
+    }
+
+    int state = best_final_state(current_scores);
+    backpointer_buffer.assign(static_cast<size_t>(reconstruction_batch_size) * STATE_COUNT, 0);
+    // Recompute one window at a time and immediately backtrace it.
+    for (int64_t window = window_count - 1; window >= 0; window--) {
+        int64_t segment_start = window * reconstruction_batch_size;
+        int64_t segment_end = segment_start + reconstruction_batch_size;
+        if (segment_end > record_length - 1) segment_end = record_length - 1;
+        run_viterbi_forward(
+            probability_bits, bases, transition_probabilities, record_length,
+            static_cast<int>(segment_start), static_cast<int>(segment_end), window == 0,
+            gt_ag_intron_score, gc_ag_intron_score, at_ac_intron_score,
+            max_boundary_intron_length,
+            stop_taa, stop_tag, stop_tga,
+            backpointer_buffer.data(), nullptr,
+            &checkpoints[static_cast<size_t>(window) * STATE_COUNT]
+        );
+        for (int64_t position = segment_end; position >= segment_start + 1; position--) {
+            path[position] = static_cast<int8_t>(state);
+            state = static_cast<int>(backpointer_buffer[
+                static_cast<size_t>(position - segment_start - 1) * STATE_COUNT + state
+            ]);
+        }
+    }
+    path[0] = static_cast<int8_t>(state);
+}
+
+extern "C" void decode_batch_on_cpu(
+    const uint16_t* probability_bits, const int8_t* bases,
+    const StepScore* transition_probabilities,
+    const int64_t* record_offsets, int record_count,
+    int stop_taa, int stop_tag, int stop_tga,
+    double gt_ag_intron_log_prior, double gc_ag_intron_log_prior,
+    double at_ac_intron_log_prior,
+    int max_boundary_intron_length,
+    int64_t requested_reconstruction_batch_size, int8_t* output_paths
+) {
+    static const bool score_lookup_ready = build_log_score_lookup();
+    (void)score_lookup_ready;
+
+    const StepScore gt_ag_intron_score = score_from_nats(gt_ag_intron_log_prior);
+    const StepScore gc_ag_intron_score = score_from_nats(gc_ag_intron_log_prior);
+    const StepScore at_ac_intron_score = score_from_nats(at_ac_intron_log_prior);
+    int64_t reconstruction_batch_size = requested_reconstruction_batch_size == 0
+        ? DEFAULT_RECONSTRUCTION_BATCH_SIZE
+        : requested_reconstruction_batch_size;
+
+#pragma omp parallel
+    {
+        std::vector<int8_t> backpointer_buffer;
+        std::vector<AccumulatedScore> checkpoints;
+#pragma omp for schedule(dynamic, 1)
+        for (int record = 0; record < record_count; record++) {
+            int64_t offset = record_offsets[record];
+            int record_length = static_cast<int>(record_offsets[record + 1] - offset);
+            decode_one_track(
+                probability_bits + offset, bases + offset,
+                transition_probabilities + static_cast<size_t>(offset) * TRANSITION_COUNT,
+                record_length,
+                gt_ag_intron_score, gc_ag_intron_score, at_ac_intron_score,
+                max_boundary_intron_length,
+                stop_taa != 0, stop_tag != 0, stop_tga != 0,
+                reconstruction_batch_size, output_paths + offset,
+                backpointer_buffer, checkpoints
+            );
+        }
+    }
+}
+
+#ifdef TORCH_EXTENSION_NAME
+#include <torch/extension.h>
+#include <pybind11/numpy.h>
+namespace py = pybind11;
+
+static py::array_t<int8_t> decode_packed_paths(
+    py::array_t<unsigned short> probability_array, py::array_t<int8_t> encoded_bases,
+    py::array_t<StepScore> transition_probabilities,
+    py::array_t<int64_t> record_offsets_array,
+    int stop_taa, int stop_tag, int stop_tga,
+    double gt_ag_intron_log_prior, double gc_ag_intron_log_prior,
+    double at_ac_intron_log_prior,
+    int max_boundary_intron_length,
+    int64_t reconstruction_batch_size
+) {
+    if (transition_probabilities.size() != encoded_bases.size() * TRANSITION_COUNT) {
+        throw std::invalid_argument(
+            "transition probabilities must contain four values per encoded base"
+        );
+    }
+    int record_count = static_cast<int>(record_offsets_array.size()) - 1;
+    auto state_paths = py::array_t<int8_t>(encoded_bases.size());
+    const uint16_t* probability_bits = probability_array.data();
+    const int8_t* base_data = encoded_bases.data();
+    const StepScore* transition_probability_data = transition_probabilities.data();
+    const int64_t* record_offsets = record_offsets_array.data();
+    int8_t* output_paths = state_paths.mutable_data();
+    {
+        // Let the Python reader prepare the next input chunk during CPU work.
+        py::gil_scoped_release release;
+        decode_batch_on_cpu(
+            probability_bits, base_data, transition_probability_data, record_offsets, record_count,
+            stop_taa, stop_tag, stop_tga,
+            gt_ag_intron_log_prior, gc_ag_intron_log_prior, at_ac_intron_log_prior,
+            max_boundary_intron_length,
+            reconstruction_batch_size, output_paths
+        );
+    }
+    return state_paths;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("decode", &decode_packed_paths, "Decode a batch with the 30-state CDS decoder");
+}
+#endif

annotation_pipeline/src/viterbi_backends/viterbi_gpu.cu

@@ -0,0 +1,1607 @@
+// viterbi_gpu.cu -- GPU Viterbi backend: CUDA implementation of the algorithm in viterbi.py.
+//
+// The active decoder has 30 states. Backpointer columns use a 32-byte stride so
+// one warp can map one lane to one state while lanes 30..31 remain padding.
+// The kernels use the same transition grammar as viterbi.py:
+//   - exon states carry one of six stop-prefix types;
+//   - GT/GC-AG introns enter FIRST, then move to READY before AG exit;
+//   - AT-AC introns use one state table and exit at AC;
+//   - Python supplies explicit GT-AG, GC-AG, and AT-AC log priors;
+//   - N freely matches any concrete base in codons and splice signals;
+//   - dedicated leading/trailing states represent introns crossing a record edge;
+//   - stop read-through is disabled.
+//
+// Transition probabilities: the transition network supplies, per base, the
+// probabilities of the four permitted transitions (enter gene, leave gene,
+// enter intron, leave intron) as log-odds against staying in the decoder's
+// fixed-point units. The grammar still decides which state changes are legal;
+// the transition probabilities re-rank the legal paths.
+//
+// Decoder stages:
+//   1. Convert probabilities to fixed-point coding/non-coding log scores.
+//   2. Split records into bounded segments. For each possible entry state,
+//      build an int32 matrix containing the best score to every exit state.
+//   3. Combine segment matrices in int64 and save the exact 30-state score
+//      vector at every segment entry while selecting each record's end state.
+//   4. Propagate each segment's entry origins with those global scores, then
+//      trace the selected segment boundaries.
+//   5. Restart each selected segment path at score 0, recompute int32
+//      backpointers, and write the final path in bounded memory windows.
+#include <cstdint>
+#include <cmath>
+#include <initializer_list>
+#include <stdexcept>
+#include <string>
+#include <vector>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+
+using StateScore = int32_t;
+using AccumulatedScore = int64_t;
+static constexpr double SCORE_SCALE = 16384.0;
+// A segment starts at zero and is short enough that every raw step score
+// remains above this int32 sentinel. Only scores accumulated across segment
+// boundaries need int64; reconstruction restarts the chosen segment path at 0.
+static constexpr StateScore UNREACHABLE_STATE_SCORE = -2000000000;
+static constexpr AccumulatedScore UNREACHABLE_ACCUMULATED_SCORE = -4000000000000000000LL;
+static constexpr int64_t RAW_SEGMENT_SCORE_LIMIT = 1800000000LL;
+
+__device__ __host__ __forceinline__ StateScore score_from_nats(double nats) {
+    return static_cast<StateScore>(llround(nats * SCORE_SCALE));
+}
+
+// Set during decoder initialization and used on the host to cap segment length
+// before any int32 segment kernel is launched.
+static int64_t maximum_absolute_step_score = 1;
+
+static const int STATE_COUNT = 30;
+static const int BACKPOINTER_STRIDE = 32;
+static const int PREFIX_TYPE_COUNT = 6;
+static const int INTERGENIC = 0;
+static const int START_A = 1;
+static const int START_AT = 2;
+static const int EXON_PREFIX_TYPE_0 = 3;
+static const int STOP = 9;
+static const int GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0 = 10;
+static const int GT_GC_AG_INTRON_READY_PREFIX_TYPE_0 = 16;
+static const int AT_AC_INTRON_PREFIX_TYPE_0 = 22;
+static const int LEADING_BOUNDARY_INTRON = 28;
+static const int TRAILING_BOUNDARY_INTRON = 29;
+
+// Transition probability tracks per base, in the column order of constants.py.
+static const int ENTER_GENE = 0;
+static const int LEAVE_GENE = 1;
+static const int ENTER_INTRON = 2;
+static const int LEAVE_INTRON = 3;
+static const int TRANSITION_COUNT = 4;
+
+__device__ __host__ __forceinline__ size_t segment_transition_index(
+    int segment, int exit_state, int entry_state
+) {
+    return (static_cast<size_t>(segment) * STATE_COUNT + exit_state) * STATE_COUNT
+        + entry_state;
+}
+
+__device__ __host__ __forceinline__ size_t segment_state_index(int segment, int state) {
+    return static_cast<size_t>(segment) * STATE_COUNT + state;
+}
+
+// State layout:
+//   0       intergenic
+//   1..2    first one/two bases of an ATG start codon
+//   3..8    exon states carrying prefix type 0..5
+//   9       completed in-frame stop codon
+//   10..15  first GT/GC-AG intron base, carrying prefix type 0..5
+//   16..21  GT/GC-AG intron ready to accept AG, carrying prefix type 0..5
+//   22..27  AT-AC intron, carrying prefix type 0..5
+//   28      intron extending from the left record boundary
+//   29      intron extending to the right record boundary
+
+static const int BASE_A = 0;
+static const int BASE_T = 1;
+static const int BASE_C = 2;
+static const int BASE_G = 3;
+static const int BASE_N = 4;
+
+__device__ __forceinline__ bool base_matches(
+    int observed_base, int expected_base
+) {
+    return observed_base == expected_base || observed_base == BASE_N;
+}
+
+__constant__ StateScore d_gt_ag_intron_score;
+__constant__ StateScore d_gc_ag_intron_score;
+__constant__ StateScore d_at_ac_intron_score;
+__constant__ int d_stop_taa_enabled;
+__constant__ int d_stop_tag_enabled;
+__constant__ int d_stop_tga_enabled;
+__constant__ int d_max_boundary_intron_length;
+__device__ __forceinline__ bool is_exact_enabled_stop_codon(
+    int first_base, int second_base, int third_base
+) {
+    if (first_base != BASE_T) return false;
+    if (second_base == BASE_A) {
+        if (third_base == BASE_A && d_stop_taa_enabled) return true;
+        if (third_base == BASE_G && d_stop_tag_enabled) return true;
+    }
+    return second_base == BASE_G
+        && third_base == BASE_A
+        && d_stop_tga_enabled;
+}
+
+__device__ __forceinline__ int exon_state_for_prefix(int prefix_type) {
+    return EXON_PREFIX_TYPE_0 + prefix_type;
+}
+
+struct PositionSignals {
+    int base;
+    StateScore coding_score;
+    StateScore noncoding_score;
+    bool has_gt_ag_donor;
+    bool has_gc_ag_donor;
+    bool has_at_ac_donor;
+    bool has_ag_acceptor;
+    bool has_ac_acceptor;
+    bool has_atg_start;
+    bool has_exact_acceptor_followed_by_atg;
+    bool has_exact_enabled_stop_followed_by_donor;
+    bool can_extend_leading_boundary_intron;
+    bool can_enter_trailing_boundary_intron;
+    bool completes_stop_from_prefix_type_3;
+    bool completes_stop_from_prefix_type_4;
+    StateScore enter_gene_log_odds;
+    StateScore leave_gene_log_odds;
+    StateScore enter_intron_log_odds;
+    StateScore leave_intron_log_odds;
+};
+
+__device__ __forceinline__ bool is_exon_state(int state) {
+    return state >= EXON_PREFIX_TYPE_0
+        && state < EXON_PREFIX_TYPE_0 + PREFIX_TYPE_COUNT;
+}
+
+__device__ __forceinline__ bool is_gt_gc_intron_entry_state(int state) {
+    return state >= GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0
+        && state < GT_GC_AG_INTRON_READY_PREFIX_TYPE_0;
+}
+
+__device__ __forceinline__ bool is_gt_gc_acceptor_state(int state) {
+    return state >= GT_GC_AG_INTRON_READY_PREFIX_TYPE_0
+        && state < AT_AC_INTRON_PREFIX_TYPE_0;
+}
+
+__device__ __forceinline__ bool is_at_ac_intron_state(int state) {
+    return state >= AT_AC_INTRON_PREFIX_TYPE_0
+        && state < LEADING_BOUNDARY_INTRON;
+}
+
+// Transition log-odds of one legal state change. Each value is attached to the
+// base where the new segment begins: the first start-codon base (enter gene),
+// the last stop-codon base (leave gene, charged on the base that closes the
+// gene), the first intron base (enter intron), or the first exon base after an
+// intron (leave intron). Introns that cross a record edge are entered and left
+// the same way.
+__device__ __forceinline__ StateScore transition_log_odds(
+    int source_state, int target_state, const PositionSignals& signals
+) {
+    if (source_state == INTERGENIC && target_state == START_A) {
+        return signals.enter_gene_log_odds;
+    }
+    // Leaving the gene takes precedence when an intron exit completes a stop
+    // codon that was split by the intron.
+    if (target_state == STOP) return signals.leave_gene_log_odds;
+    if (
+        is_exon_state(source_state)
+        && (is_gt_gc_intron_entry_state(target_state)
+            || is_at_ac_intron_state(target_state)
+            || target_state == TRAILING_BOUNDARY_INTRON)
+    ) {
+        return signals.enter_intron_log_odds;
+    }
+    if (
+        (is_gt_gc_acceptor_state(source_state)
+            || is_at_ac_intron_state(source_state)
+            || source_state == LEADING_BOUNDARY_INTRON)
+        && is_exon_state(target_state)
+    ) {
+        return signals.leave_intron_log_odds;
+    }
+    return StateScore(0);
+}
+
+__device__ __forceinline__ StateScore best_gt_gc_ag_intron_score(
+    const PositionSignals& signals
+) {
+    StateScore best_score = UNREACHABLE_STATE_SCORE;
+    if (signals.has_gt_ag_donor) best_score = d_gt_ag_intron_score;
+    if (signals.has_gc_ag_donor && d_gc_ag_intron_score > best_score) {
+        best_score = d_gc_ag_intron_score;
+    }
+    return best_score;
+}
+
+__device__ __forceinline__ StateScore best_intron_score_for_donor(
+    const PositionSignals& signals
+) {
+    StateScore best_score = best_gt_gc_ag_intron_score(signals);
+    if (signals.has_at_ac_donor && d_at_ac_intron_score > best_score) {
+        best_score = d_at_ac_intron_score;
+    }
+    return best_score;
+}
+
+template <typename Score>
+struct TransitionChoice {
+    Score score;
+    int source_state;
+};
+
+__device__ __forceinline__ bool is_unreachable(StateScore score) {
+    return score == UNREACHABLE_STATE_SCORE;
+}
+
+__device__ __forceinline__ bool is_unreachable(AccumulatedScore score) {
+    return score == UNREACHABLE_ACCUMULATED_SCORE;
+}
+
+__device__ __forceinline__ StateScore unreachable_score(StateScore) {
+    return UNREACHABLE_STATE_SCORE;
+}
+
+__device__ __forceinline__ AccumulatedScore unreachable_score(AccumulatedScore) {
+    return UNREACHABLE_ACCUMULATED_SCORE;
+}
+
+__device__ __forceinline__ int next_coding_state_after_base(
+    int prefix_type, int base,
+    bool completes_stop_from_prefix_type_3,
+    bool completes_stop_from_prefix_type_4
+) {
+    if (prefix_type == 0) return exon_state_for_prefix(base == BASE_T ? 1 : 2);
+    if (prefix_type == 1) {
+        if (base == BASE_A) return exon_state_for_prefix(3);
+        if (base == BASE_G) return exon_state_for_prefix(4);
+        return exon_state_for_prefix(5);
+    }
+    if (prefix_type == 2) return exon_state_for_prefix(5);
+    if (prefix_type == 3) {
+        if (completes_stop_from_prefix_type_3) return STOP;
+        return exon_state_for_prefix(0);
+    }
+    if (prefix_type == 4) {
+        if (completes_stop_from_prefix_type_4) return STOP;
+        return exon_state_for_prefix(0);
+    }
+    return exon_state_for_prefix(0);
+}
+
+__device__ __forceinline__ int next_coding_state_after_base(
+    int prefix_type, const PositionSignals& signals
+) {
+    return next_coding_state_after_base(
+        prefix_type, signals.base,
+        signals.completes_stop_from_prefix_type_3, signals.completes_stop_from_prefix_type_4
+    );
+}
+
+__device__ __forceinline__ int next_coding_state_after_resolving_n(
+    int prefix_type, int resolved_base
+) {
+    bool completes_stop_from_prefix_type_3 =
+        (resolved_base == BASE_A && d_stop_taa_enabled)
+        || (resolved_base == BASE_G && d_stop_tag_enabled);
+    bool completes_stop_from_prefix_type_4 =
+        resolved_base == BASE_A && d_stop_tga_enabled;
+    return next_coding_state_after_base(
+        prefix_type, resolved_base,
+        completes_stop_from_prefix_type_3, completes_stop_from_prefix_type_4
+    );
+}
+
+__device__ __forceinline__ bool coding_base_can_reach_state(
+    int prefix_type, int target_state, const PositionSignals& signals
+) {
+    if (signals.base != BASE_N) {
+        return target_state == next_coding_state_after_base(prefix_type, signals);
+    }
+    for (int resolved_base = BASE_A; resolved_base <= BASE_G; resolved_base++) {
+        if (target_state == next_coding_state_after_resolving_n(
+            prefix_type, resolved_base
+        )) {
+            return true;
+        }
+    }
+    return false;
+}
+
+__device__ __forceinline__ PositionSignals compute_position_signals(
+    const StateScore* coding_scores, const StateScore* noncoding_scores,
+    const StateScore* transition_probabilities, const int8_t* bases,
+    int record_length, int position
+) {
+    // position is always at least 1. At a record edge, unavailable motif bases
+    // are unknown rather than evidence that a start, donor, or acceptor is
+    // impossible. Internal positions still require the complete motif.
+    int base = bases[position];
+    bool next_base_is_missing = position + 1 == record_length;
+    int next_base = next_base_is_missing ? -1 : bases[position + 1];
+    bool has_gt_ag_donor = base_matches(base, BASE_G)
+        && (next_base_is_missing || base_matches(next_base, BASE_T));
+    bool has_gc_ag_donor = base_matches(base, BASE_G)
+        && (next_base_is_missing || base_matches(next_base, BASE_C));
+    bool has_at_ac_donor = base_matches(base, BASE_A)
+        && (next_base_is_missing || base_matches(next_base, BASE_T));
+    bool acceptor_first_base_matches = position == 1
+        || base_matches(bases[position - 2], BASE_A);
+    bool has_ag_acceptor = acceptor_first_base_matches
+        && base_matches(bases[position - 1], BASE_G);
+    bool has_ac_acceptor = acceptor_first_base_matches
+        && base_matches(bases[position - 1], BASE_C);
+    bool third_start_base_matches = position + 2 >= record_length
+        || base_matches(bases[position + 2], BASE_G);
+    bool has_atg_start = base_matches(base, BASE_A)
+        && (next_base_is_missing || base_matches(next_base, BASE_T))
+        && third_start_base_matches;
+    // Boundary guards require exact five-base patterns. N and missing edge
+    // context remain valid for ordinary motif matching only.
+    bool has_exact_acceptor_followed_by_atg = position >= 2
+        && position + 2 < record_length
+        && bases[position - 2] == BASE_A
+        && (bases[position - 1] == BASE_G || bases[position - 1] == BASE_C)
+        && base == BASE_A
+        && next_base == BASE_T
+        && bases[position + 2] == BASE_G;
+    bool has_exact_donor = position + 1 < record_length
+        && ((base == BASE_G && (next_base == BASE_T || next_base == BASE_C))
+            || (base == BASE_A && next_base == BASE_T));
+    bool has_exact_enabled_stop_followed_by_donor = false;
+    if (has_exact_donor && position >= 3) {
+        has_exact_enabled_stop_followed_by_donor = is_exact_enabled_stop_codon(
+            bases[position - 3], bases[position - 2], bases[position - 1]
+        );
+    }
+    bool completes_stop_from_prefix_type_3 =
+        (base_matches(base, BASE_A) && d_stop_taa_enabled)
+        || (base_matches(base, BASE_G) && d_stop_tag_enabled);
+    bool completes_stop_from_prefix_type_4 =
+        base_matches(base, BASE_A) && d_stop_tga_enabled;
+    // A leading self-loop into position makes its run position + 1 bases.
+    // A trailing entry here emits record_length - position bases through the
+    // right edge, so neither boundary needs a counter state.
+    bool boundary_introns_enabled = d_max_boundary_intron_length != 0;
+    bool boundary_intron_length_is_unlimited = d_max_boundary_intron_length == -1;
+    bool can_extend_leading_boundary_intron =
+        boundary_introns_enabled
+        && (boundary_intron_length_is_unlimited
+            || position < d_max_boundary_intron_length);
+    bool can_enter_trailing_boundary_intron =
+        boundary_introns_enabled
+        && (boundary_intron_length_is_unlimited
+            || record_length - position <= d_max_boundary_intron_length);
+    const StateScore* position_scores =
+        transition_probabilities + static_cast<size_t>(position) * TRANSITION_COUNT;
+    return {
+        base, coding_scores[position], noncoding_scores[position],
+        has_gt_ag_donor, has_gc_ag_donor, has_at_ac_donor,
+        has_ag_acceptor, has_ac_acceptor, has_atg_start,
+        has_exact_acceptor_followed_by_atg,
+        has_exact_enabled_stop_followed_by_donor,
+        can_extend_leading_boundary_intron,
+        can_enter_trailing_boundary_intron,
+        completes_stop_from_prefix_type_3, completes_stop_from_prefix_type_4,
+        position_scores[ENTER_GENE],
+        position_scores[LEAVE_GENE],
+        position_scores[ENTER_INTRON],
+        position_scores[LEAVE_INTRON]
+    };
+}
+
+__device__ __forceinline__ void keep_best_transition(
+    StateScore* next_scores, int8_t* predecessors,
+    int target_state, StateScore candidate_score, int source_state,
+    const PositionSignals& signals
+) {
+    candidate_score += transition_log_odds(
+        source_state, target_state, signals
+    );
+    if (candidate_score > next_scores[target_state]) {
+        next_scores[target_state] = candidate_score;
+        if (predecessors) predecessors[target_state] = static_cast<int8_t>(source_state);
+    }
+}
+
+__device__ __forceinline__ void emit_coding_base(
+    const StateScore* current_scores, StateScore* next_scores, int8_t* predecessors,
+    int source_state, int prefix_type, StateScore coding_score,
+    const PositionSignals& signals
+) {
+    StateScore source_score = current_scores[source_state];
+    if (source_score == UNREACHABLE_STATE_SCORE) return;
+    StateScore candidate_score = source_score + coding_score;
+
+    if (signals.base == BASE_N) {
+        for (int resolved_base = BASE_A; resolved_base <= BASE_G; resolved_base++) {
+            int target_state = next_coding_state_after_resolving_n(
+                prefix_type, resolved_base
+            );
+            keep_best_transition(
+                next_scores, predecessors, target_state, candidate_score,
+                source_state, signals
+            );
+        }
+    } else {
+        int target_state = next_coding_state_after_base(prefix_type, signals);
+        keep_best_transition(
+            next_scores, predecessors, target_state, candidate_score,
+            source_state, signals
+        );
+    }
+}
+
+template <typename Score>
+__device__ __forceinline__ void consider_predecessor(
+    TransitionChoice<Score>& best, const Score* current_scores,
+    int source_state, int target_state, StateScore added_score,
+    const PositionSignals& signals
+) {
+    Score source_score = current_scores[source_state];
+    if (is_unreachable(source_score)) return;
+
+    Score candidate_score = source_score + static_cast<Score>(
+        added_score + transition_log_odds(
+            source_state, target_state, signals
+        )
+    );
+    if (candidate_score > best.score) {
+        best.score = candidate_score;
+        best.source_state = source_state;
+    }
+}
+
+// A warp assigns one target state to each lane. Gathering only that target's
+// legal predecessors prevents every lane from repeating all source updates.
+template <typename Score>
+__device__ __forceinline__ TransitionChoice<Score> best_transition_to_state(
+    int target_state, const Score* current_scores,
+    const PositionSignals& signals
+) {
+    TransitionChoice<Score> best = {unreachable_score(Score{}), 0};
+
+    if (target_state == INTERGENIC) {
+        consider_predecessor(
+            best, current_scores, INTERGENIC, target_state,
+            signals.noncoding_score, signals
+        );
+        consider_predecessor(
+            best, current_scores, STOP, target_state,
+            signals.noncoding_score, signals
+        );
+        return best;
+    }
+    if (target_state == START_A) {
+        if (signals.has_atg_start) {
+            consider_predecessor(
+                best, current_scores, INTERGENIC, target_state,
+                signals.coding_score, signals
+            );
+        }
+        return best;
+    }
+    if (target_state == START_AT) {
+        consider_predecessor(
+            best, current_scores, START_A, target_state,
+            signals.coding_score, signals
+        );
+        return best;
+    }
+
+    bool is_exon_or_stop = target_state == STOP || (
+        target_state >= EXON_PREFIX_TYPE_0
+        && target_state < EXON_PREFIX_TYPE_0 + PREFIX_TYPE_COUNT
+    );
+    if (is_exon_or_stop) {
+        if (target_state == exon_state_for_prefix(0)) {
+            consider_predecessor(
+                best, current_scores, START_AT, target_state,
+                signals.coding_score, signals
+            );
+        }
+        // Every prefix type is allowed when the leading boundary intron exits,
+        // so the boundary state does not need to carry one.
+        if ((signals.has_ag_acceptor || signals.has_ac_acceptor)
+            && !signals.has_exact_acceptor_followed_by_atg) {
+#pragma unroll
+            for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+                if (coding_base_can_reach_state(prefix_type, target_state, signals)) {
+                    consider_predecessor(
+                        best, current_scores, LEADING_BOUNDARY_INTRON,
+                        target_state, signals.coding_score,
+                        signals
+                    );
+                }
+            }
+        }
+#pragma unroll
+        for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+            if (coding_base_can_reach_state(prefix_type, target_state, signals)) {
+                consider_predecessor(
+                    best, current_scores, exon_state_for_prefix(prefix_type),
+                    target_state, signals.coding_score, signals
+                );
+            }
+        }
+        if (signals.has_ag_acceptor) {
+#pragma unroll
+            for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+                if (coding_base_can_reach_state(prefix_type, target_state, signals)) {
+                    consider_predecessor(
+                        best, current_scores,
+                        GT_GC_AG_INTRON_READY_PREFIX_TYPE_0 + prefix_type,
+                        target_state, signals.coding_score,
+                        signals
+                    );
+                }
+            }
+        }
+        if (signals.has_ac_acceptor) {
+#pragma unroll
+            for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+                if (coding_base_can_reach_state(prefix_type, target_state, signals)) {
+                    consider_predecessor(
+                        best, current_scores, AT_AC_INTRON_PREFIX_TYPE_0 + prefix_type,
+                        target_state, signals.coding_score,
+                        signals
+                    );
+                }
+            }
+        }
+        return best;
+    }
+
+    if (target_state >= GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0
+        && target_state < GT_GC_AG_INTRON_READY_PREFIX_TYPE_0) {
+        if (signals.has_gt_ag_donor || signals.has_gc_ag_donor) {
+            int prefix_type = target_state - GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0;
+            StateScore intron_score = best_gt_gc_ag_intron_score(signals);
+            consider_predecessor(
+                best, current_scores, exon_state_for_prefix(prefix_type),
+                target_state,
+                signals.noncoding_score + intron_score,
+                signals
+            );
+        }
+        return best;
+    }
+
+    if (target_state >= GT_GC_AG_INTRON_READY_PREFIX_TYPE_0
+        && target_state < AT_AC_INTRON_PREFIX_TYPE_0) {
+        int prefix_type = target_state - GT_GC_AG_INTRON_READY_PREFIX_TYPE_0;
+        consider_predecessor(
+            best, current_scores, GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0 + prefix_type,
+            target_state, signals.noncoding_score, signals
+        );
+        consider_predecessor(
+            best, current_scores, target_state, target_state,
+            signals.noncoding_score, signals
+        );
+        return best;
+    }
+
+    if (target_state >= AT_AC_INTRON_PREFIX_TYPE_0
+        && target_state < LEADING_BOUNDARY_INTRON) {
+        int prefix_type = target_state - AT_AC_INTRON_PREFIX_TYPE_0;
+        if (signals.has_at_ac_donor) {
+            consider_predecessor(
+                best, current_scores, exon_state_for_prefix(prefix_type),
+                target_state,
+                signals.noncoding_score + d_at_ac_intron_score,
+                signals
+            );
+        }
+        consider_predecessor(
+            best, current_scores, target_state, target_state,
+            signals.noncoding_score, signals
+        );
+        return best;
+    }
+
+    if (target_state == LEADING_BOUNDARY_INTRON) {
+        if (signals.can_extend_leading_boundary_intron) {
+            consider_predecessor(
+                best, current_scores, LEADING_BOUNDARY_INTRON,
+                target_state, signals.noncoding_score, signals
+            );
+        }
+        return best;
+    }
+
+    if (target_state == TRAILING_BOUNDARY_INTRON) {
+        bool has_any_donor = signals.has_gt_ag_donor
+            || signals.has_gc_ag_donor
+            || signals.has_at_ac_donor;
+        if (has_any_donor
+            && !signals.has_exact_enabled_stop_followed_by_donor
+            && signals.can_enter_trailing_boundary_intron) {
+            StateScore intron_score = best_intron_score_for_donor(signals);
+            StateScore added_score =
+                signals.noncoding_score + intron_score;
+#pragma unroll
+            for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+                consider_predecessor(
+                    best, current_scores, exon_state_for_prefix(prefix_type),
+                    target_state, added_score, signals
+                );
+            }
+        }
+        consider_predecessor(
+            best, current_scores, TRAILING_BOUNDARY_INTRON, target_state,
+            signals.noncoding_score, signals
+        );
+        return best;
+    }
+
+    return best;
+}
+
+// A scalar matrix thread owns all 30 target scores, so it is faster and clearer
+// to brush transitions from each reachable source state. The warp path above
+// uses target-oriented gathering because each lane owns only one target.
+__device__ inline void advance_one_column(
+    const StateScore* current_scores, StateScore* next_scores, int8_t* predecessors,
+    const StateScore* coding_scores, const StateScore* noncoding_scores,
+    const StateScore* transition_probabilities, const int8_t* bases,
+    int record_length, int position
+) {
+    PositionSignals signals = compute_position_signals(
+        coding_scores, noncoding_scores, transition_probabilities,
+        bases, record_length, position
+    );
+
+#pragma unroll
+    for (int state = 0; state < STATE_COUNT; state++) {
+        next_scores[state] = UNREACHABLE_STATE_SCORE;
+        if (predecessors) predecessors[state] = 0;
+    }
+
+    // Intergenic, stop, and start-codon transitions.
+    StateScore source_score = current_scores[INTERGENIC];
+    if (source_score != UNREACHABLE_STATE_SCORE) {
+        keep_best_transition(
+            next_scores, predecessors, INTERGENIC,
+            source_score + signals.noncoding_score, INTERGENIC, signals
+        );
+        if (signals.has_atg_start) {
+            keep_best_transition(
+                next_scores, predecessors, START_A,
+                source_score + signals.coding_score,
+                INTERGENIC, signals
+            );
+        }
+    }
+
+    source_score = current_scores[STOP];
+    if (source_score != UNREACHABLE_STATE_SCORE) {
+        keep_best_transition(
+            next_scores, predecessors, INTERGENIC,
+            source_score + signals.noncoding_score,
+            STOP, signals
+        );
+    }
+
+    source_score = current_scores[START_A];
+    if (source_score != UNREACHABLE_STATE_SCORE) {
+        keep_best_transition(
+            next_scores, predecessors, START_AT,
+            source_score + signals.coding_score, START_A, signals
+        );
+    }
+    source_score = current_scores[START_AT];
+    if (source_score != UNREACHABLE_STATE_SCORE) {
+        keep_best_transition(
+            next_scores, predecessors, exon_state_for_prefix(0),
+            source_score + signals.coding_score, START_AT, signals
+        );
+    }
+
+    // Every prefix type is allowed when the leading boundary intron exits, so
+    // the state carries none. Only an exact acceptor followed by an exact ATG
+    // activates the guard; N leaves both paths available.
+    source_score = current_scores[LEADING_BOUNDARY_INTRON];
+    if (source_score != UNREACHABLE_STATE_SCORE) {
+        if (signals.can_extend_leading_boundary_intron) {
+            keep_best_transition(
+                next_scores, predecessors,
+                LEADING_BOUNDARY_INTRON, source_score + signals.noncoding_score,
+                LEADING_BOUNDARY_INTRON, signals
+            );
+        }
+        bool can_exit_leading_intron =
+            (signals.has_ag_acceptor || signals.has_ac_acceptor)
+            && !signals.has_exact_acceptor_followed_by_atg;
+        if (can_exit_leading_intron) {
+#pragma unroll
+            for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+                emit_coding_base(
+                    current_scores, next_scores, predecessors,
+                    LEADING_BOUNDARY_INTRON, prefix_type,
+                    signals.coding_score,
+                    signals
+                );
+            }
+        }
+    }
+
+    // Exons emit coding bases, open an ordinary internal intron, or enter the
+    // trailing state when an intron continues beyond the record.
+    bool has_gt_gc_ag_donor = signals.has_gt_ag_donor || signals.has_gc_ag_donor;
+    StateScore gt_gc_ag_intron_score = best_gt_gc_ag_intron_score(signals);
+    StateScore any_intron_score = best_intron_score_for_donor(signals);
+    bool can_enter_trailing =
+        (has_gt_gc_ag_donor || signals.has_at_ac_donor)
+        && !signals.has_exact_enabled_stop_followed_by_donor
+        && signals.can_enter_trailing_boundary_intron;
+#pragma unroll
+    for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+        int exon_state = exon_state_for_prefix(prefix_type);
+        source_score = current_scores[exon_state];
+        if (source_score == UNREACHABLE_STATE_SCORE) continue;
+
+        emit_coding_base(
+            current_scores, next_scores, predecessors,
+            exon_state, prefix_type, signals.coding_score, signals
+        );
+        if (has_gt_gc_ag_donor) {
+            keep_best_transition(
+                next_scores, predecessors,
+                GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0 + prefix_type,
+                source_score + signals.noncoding_score + gt_gc_ag_intron_score,
+                exon_state, signals
+            );
+        }
+        if (signals.has_at_ac_donor) {
+            keep_best_transition(
+                next_scores, predecessors,
+                AT_AC_INTRON_PREFIX_TYPE_0 + prefix_type,
+                source_score + signals.noncoding_score + d_at_ac_intron_score,
+                exon_state, signals
+            );
+        }
+        if (can_enter_trailing) {
+            keep_best_transition(
+                next_scores, predecessors,
+                TRAILING_BOUNDARY_INTRON,
+                source_score + signals.noncoding_score + any_intron_score,
+                exon_state, signals
+            );
+        }
+    }
+
+    // FIRST must consume one additional intron base before AG exit is legal.
+#pragma unroll
+    for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+        int first_state = GT_GC_AG_INTRON_FIRST_PREFIX_TYPE_0 + prefix_type;
+        source_score = current_scores[first_state];
+        if (source_score != UNREACHABLE_STATE_SCORE) {
+            keep_best_transition(
+                next_scores, predecessors,
+                GT_GC_AG_INTRON_READY_PREFIX_TYPE_0 + prefix_type,
+                source_score + signals.noncoding_score,
+                first_state, signals
+            );
+        }
+    }
+
+    // A trailing boundary intron cannot return to an exon.
+    source_score = current_scores[TRAILING_BOUNDARY_INTRON];
+    if (source_score != UNREACHABLE_STATE_SCORE) {
+        keep_best_transition(
+            next_scores, predecessors,
+            TRAILING_BOUNDARY_INTRON, source_score + signals.noncoding_score,
+            TRAILING_BOUNDARY_INTRON, signals
+        );
+    }
+
+    // READY self-loops or exits a GT/GC-AG intron after an AG acceptor.
+#pragma unroll
+    for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+        int intron_state = GT_GC_AG_INTRON_READY_PREFIX_TYPE_0 + prefix_type;
+        source_score = current_scores[intron_state];
+        if (source_score == UNREACHABLE_STATE_SCORE) continue;
+
+        keep_best_transition(
+            next_scores, predecessors, intron_state,
+            source_score + signals.noncoding_score, intron_state, signals
+        );
+        if (signals.has_ag_acceptor) {
+            emit_coding_base(
+                current_scores, next_scores, predecessors,
+                intron_state, prefix_type, signals.coding_score,
+                signals
+            );
+        }
+    }
+
+    // AT-AC introns self-loop or exit after AC; no FIRST state is required.
+#pragma unroll
+    for (int prefix_type = 0; prefix_type < PREFIX_TYPE_COUNT; prefix_type++) {
+        int intron_state = AT_AC_INTRON_PREFIX_TYPE_0 + prefix_type;
+        source_score = current_scores[intron_state];
+        if (source_score == UNREACHABLE_STATE_SCORE) continue;
+
+        keep_best_transition(
+            next_scores, predecessors, intron_state,
+            source_score + signals.noncoding_score, intron_state, signals
+        );
+        if (signals.has_ac_acceptor) {
+            emit_coding_base(
+                current_scores, next_scores, predecessors,
+                intron_state, prefix_type, signals.coding_score,
+                signals
+            );
+        }
+    }
+}
+
+// Transitions across the left record edge are charged like interior ones: a
+// start codon beginning at the first base is charged enter_gene[0] and a stop
+// codon completed there is charged leave_gene[0]; every other first-base state
+// is charged nothing.
+__device__ __forceinline__ StateScore first_base_score_for_state(
+    int state, const int8_t* bases, int record_length,
+    StateScore first_coding_score, StateScore first_noncoding_score,
+    StateScore first_enter_gene_score, StateScore first_leave_gene_score
+) {
+    int base = bases[0];
+    bool second_start_base_matches = record_length == 1
+        || base_matches(bases[1], BASE_T);
+    bool third_start_base_matches = record_length <= 2
+        || base_matches(bases[2], BASE_G);
+    bool has_atg_start = base_matches(base, BASE_A)
+        && second_start_base_matches
+        && third_start_base_matches;
+
+    if (state == INTERGENIC) return first_noncoding_score;
+    if (state == START_A && has_atg_start) return first_coding_score + first_enter_gene_score;
+    if (state == STOP) {
+        bool can_end_stop = (
+            base_matches(base, BASE_A)
+            && (d_stop_taa_enabled || d_stop_tga_enabled)
+        ) || (base_matches(base, BASE_G) && d_stop_tag_enabled);
+        return can_end_stop ? first_coding_score + first_leave_gene_score : UNREACHABLE_STATE_SCORE;
+    }
+    if (base == BASE_N
+        && state >= EXON_PREFIX_TYPE_0
+        && state < EXON_PREFIX_TYPE_0 + PREFIX_TYPE_COUNT) {
+        return first_coding_score;
+    }
+    if (state == exon_state_for_prefix(0) || state == exon_state_for_prefix(5)) {
+        return first_coding_score;
+    }
+    if (state == exon_state_for_prefix(1) && base == BASE_T) return first_coding_score;
+    if (state == exon_state_for_prefix(2) && base != BASE_T) return first_coding_score;
+    if (state == exon_state_for_prefix(3) && base == BASE_A) return first_coding_score;
+    if (state == exon_state_for_prefix(4) && base == BASE_G) return first_coding_score;
+    if (state == LEADING_BOUNDARY_INTRON && d_max_boundary_intron_length != 0) {
+        return first_noncoding_score;
+    }
+    return UNREACHABLE_STATE_SCORE;
+}
+
+__device__ inline void decode_segment_from_one_entry_state(
+    const StateScore* coding_scores, const StateScore* noncoding_scores,
+    const StateScore* transition_probabilities, const int8_t* bases, int record_length,
+    int segment_start, int segment_end, int entry_state,
+    StateScore* output_scores
+) {
+    StateScore current_scores[STATE_COUNT], next_scores[STATE_COUNT];
+#pragma unroll
+    for (int state = 0; state < STATE_COUNT; state++) {
+        current_scores[state] = UNREACHABLE_STATE_SCORE;
+    }
+    current_scores[entry_state] = StateScore(0);
+
+    for (int position = segment_start + 1; position <= segment_end; position++) {
+        advance_one_column(
+            current_scores, next_scores, nullptr,
+            coding_scores, noncoding_scores, transition_probabilities,
+            bases, record_length, position
+        );
+#pragma unroll
+        for (int state = 0; state < STATE_COUNT; state++) {
+            current_scores[state] = next_scores[state];
+        }
+    }
+
+#pragma unroll
+    for (int state = 0; state < STATE_COUNT; state++) {
+        output_scores[state] = current_scores[state];
+    }
+}
+
+__device__ StateScore d_coding_score_lookup[65536], d_noncoding_score_lookup[65536];
+
+__global__ void build_log_score_lookup() {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= 65536) return;
+    double p = (double)__half2float(__ushort_as_half((unsigned short)idx));
+    if (p < 1e-6) p = 1e-6;
+    else if (p > 1.0 - 1e-6) p = 1.0 - 1e-6;
+    d_coding_score_lookup[idx] = score_from_nats(log(p));
+    d_noncoding_score_lookup[idx] = score_from_nats(log(1.0 - p));
+}
+
+__global__ void convert_probabilities_to_log_scores(
+    const __half* probabilities, StateScore* coding_scores,
+    StateScore* noncoding_scores, int base_count
+) {
+    int base_index = blockIdx.x * blockDim.x + threadIdx.x;
+    if (base_index >= base_count) return;
+    unsigned short probability_bits = __half_as_ushort(probabilities[base_index]);
+    coding_scores[base_index] = d_coding_score_lookup[probability_bits];
+    noncoding_scores[base_index] = d_noncoding_score_lookup[probability_bits];
+}
+
+static StateScore configured_step_score(double nats, const char* parameter_name) {
+    double scaled_score = nats * SCORE_SCALE;
+    if (!std::isfinite(scaled_score)
+        || scaled_score <= -RAW_SEGMENT_SCORE_LIMIT
+        || scaled_score >= RAW_SEGMENT_SCORE_LIMIT) {
+        throw std::invalid_argument(
+            std::string(parameter_name) + " is outside the supported fixed-point range"
+        );
+    }
+    return static_cast<StateScore>(llround(scaled_score));
+}
+
+extern "C" void initialize_gpu_decoder(
+    int stop_taa, int stop_tag, int stop_tga,
+    double gt_ag_intron_log_prior, double gc_ag_intron_log_prior,
+    double at_ac_intron_log_prior, int max_boundary_intron_length
+) {
+    StateScore gt_ag_intron_score = configured_step_score(
+        gt_ag_intron_log_prior, "gt_ag_intron_log_prior"
+    );
+    StateScore gc_ag_intron_score = configured_step_score(
+        gc_ag_intron_log_prior, "gc_ag_intron_log_prior"
+    );
+    StateScore at_ac_intron_score = configured_step_score(
+        at_ac_intron_log_prior, "at_ac_intron_log_prior"
+    );
+    // Every step emits one base and adds at most one intron prior; transition
+    // scores are bounded separately per batch. This bound lets the host shorten
+    // risky segments before their raw int32 scores are computed.
+    StateScore minimum_emission_score = score_from_nats(log(1e-6));
+    StateScore maximum_emission_score = score_from_nats(log(1.0 - 1e-6));
+    StateScore intron_scores[] = {gt_ag_intron_score, gc_ag_intron_score, at_ac_intron_score};
+    maximum_absolute_step_score = 1;
+    auto include_step_score = [&](int64_t score) {
+        int64_t magnitude = score < 0 ? -score : score;
+        if (magnitude >= RAW_SEGMENT_SCORE_LIMIT) {
+            throw std::invalid_argument(
+                "configured priors cannot be represented in one int32 segment step"
+            );
+        }
+        if (magnitude > maximum_absolute_step_score) {
+            maximum_absolute_step_score = magnitude;
+        }
+    };
+    for (StateScore emission_score : {minimum_emission_score, maximum_emission_score}) {
+        include_step_score(emission_score);
+        for (StateScore intron_score : intron_scores) {
+            include_step_score(static_cast<int64_t>(emission_score) + intron_score);
+        }
+    }
+
+    cudaMemcpyToSymbol(d_gt_ag_intron_score, &gt_ag_intron_score, sizeof(StateScore));
+    cudaMemcpyToSymbol(d_gc_ag_intron_score, &gc_ag_intron_score, sizeof(StateScore));
+    cudaMemcpyToSymbol(d_at_ac_intron_score, &at_ac_intron_score, sizeof(StateScore));
+    cudaMemcpyToSymbol(d_stop_taa_enabled, &stop_taa, sizeof(int));
+    cudaMemcpyToSymbol(d_stop_tag_enabled, &stop_tag, sizeof(int));
+    cudaMemcpyToSymbol(d_stop_tga_enabled, &stop_tga, sizeof(int));
+    cudaMemcpyToSymbol(d_max_boundary_intron_length, &max_boundary_intron_length, sizeof(int));
+    build_log_score_lookup<<<256, 256>>>();
+}
+
+__global__ void build_segment_transition_matrices(
+    const StateScore* coding_scores, const StateScore* noncoding_scores,
+    const StateScore* transition_probabilities, const int8_t* bases,
+    const int64_t* record_offsets,
+    const int* record_lengths,
+    const int64_t* segment_starts, const int64_t* segment_ends,
+    int segment_count, StateScore* transition_matrices
+) {
+    // One task decodes a segment from one possible entry state. Its 30 final
+    // scores form one column of the segment's max-plus transition matrix.
+    int task_index = blockIdx.x * blockDim.x + threadIdx.x;
+    int segment = task_index / STATE_COUNT;
+    int entry_state = task_index % STATE_COUNT;
+    if (segment >= segment_count) return;
+
+    int64_t offset = record_offsets[segment];
+    int segment_start = static_cast<int>(segment_starts[segment] - offset);
+    int segment_end = static_cast<int>(segment_ends[segment] - offset);
+    StateScore output_scores[STATE_COUNT];
+    decode_segment_from_one_entry_state(
+        coding_scores + offset, noncoding_scores + offset,
+        transition_probabilities + offset * TRANSITION_COUNT, bases + offset,
+        record_lengths[segment], segment_start, segment_end,
+        entry_state, output_scores
+    );
+#pragma unroll
+    for (int exit_state = 0; exit_state < STATE_COUNT; exit_state++) {
+        transition_matrices[
+            segment_transition_index(segment, exit_state, entry_state)
+        ] = output_scores[exit_state];
+    }
+}
+
+__global__ void recompute_selected_segment_backpointers(
+    const StateScore* coding_scores, const StateScore* noncoding_scores,
+    const StateScore* transition_probabilities, const int8_t* bases,
+    const int64_t* segment_record_offsets,
+    const int* segment_record_lengths,
+    const int64_t* segment_starts, const int64_t* segment_ends,
+    const int8_t* segment_entry_states,
+    int segment_count, int64_t backpointer_origin,
+    int8_t* backpointers
+) {
+    int warp_in_block = threadIdx.x >> 5;
+    int segment = blockIdx.x * 4 + warp_in_block;
+    if (segment >= segment_count) return;
+
+    int target_state = threadIdx.x & 31;
+    int64_t offset = segment_record_offsets[segment];
+    int record_length = segment_record_lengths[segment];
+    int segment_start = static_cast<int>(segment_starts[segment] - offset);
+    int segment_end = static_cast<int>(segment_ends[segment] - offset);
+    const StateScore* record_coding_scores = coding_scores + offset;
+    const StateScore* record_noncoding_scores = noncoding_scores + offset;
+    const StateScore* record_transition_probabilities =
+        transition_probabilities + offset * TRANSITION_COUNT;
+    const int8_t* record_bases = bases + offset;
+
+    int entry_state = static_cast<int>(segment_entry_states[segment]);
+    __shared__ StateScore shared_scores[4][STATE_COUNT];
+    StateScore* current_scores = shared_scores[warp_in_block];
+    if (target_state < STATE_COUNT) {
+        current_scores[target_state] = target_state == entry_state
+            ? StateScore(0)
+            : UNREACHABLE_STATE_SCORE;
+    }
+    __syncwarp();
+
+    for (int position = segment_start + 1; position <= segment_end; position++) {
+        PositionSignals signals = compute_position_signals(
+            record_coding_scores, record_noncoding_scores,
+            record_transition_probabilities, record_bases, record_length, position
+        );
+        TransitionChoice<StateScore> best = {UNREACHABLE_STATE_SCORE, 0};
+        if (target_state < STATE_COUNT) {
+            best = best_transition_to_state(target_state, current_scores, signals);
+        }
+
+        if (target_state < STATE_COUNT) {
+            current_scores[target_state] = best.score;
+            backpointers[
+                static_cast<size_t>(offset + position - 1 - backpointer_origin)
+                    * BACKPOINTER_STRIDE
+                + target_state
+            ] =
+                static_cast<int8_t>(best.source_state);
+        }
+        __syncwarp();
+    }
+}
+
+__global__ void write_selected_segment_paths(
+    const int64_t* segment_starts, const int64_t* segment_ends,
+    const int* segment_exit_states, int segment_count,
+    int64_t backpointer_origin, const int8_t* backpointers,
+    int8_t* paths
+) {
+    int segment = blockIdx.x * blockDim.x + threadIdx.x;
+    if (segment >= segment_count) return;
+
+    int64_t segment_start = segment_starts[segment];
+    int64_t segment_end = segment_ends[segment];
+    int current_state = segment_exit_states[segment];
+    for (int64_t position = segment_end; position > segment_start; position--) {
+        paths[position] = static_cast<int8_t>(current_state);
+        current_state = static_cast<int>(
+            backpointers[
+                static_cast<size_t>(position - 1 - backpointer_origin) * BACKPOINTER_STRIDE
+                + current_state
+            ]
+        );
+    }
+    paths[segment_start] = static_cast<int8_t>(current_state);
+}
+
+__global__ void combine_segment_scores(
+    const StateScore* coding_scores, const StateScore* noncoding_scores,
+    const StateScore* transition_probabilities, const int8_t* bases,
+    const int64_t* record_offsets, const int* record_lengths,
+    const StateScore* transition_matrices,
+    const int* record_first_segments, const int* record_segment_counts, int record_count,
+    AccumulatedScore* segment_entry_scores, int* record_end_states
+) {
+    // One warp handles one record. Lane n owns state n and repeatedly applies
+    // the record's segment matrices. Scores crossing segment boundaries use
+    // int64, so arbitrarily long records do not overflow.
+    int record = blockIdx.x;
+    if (record >= record_count) return;
+    int state = threadIdx.x;
+    int segment_count = record_segment_counts[record];
+    int first_segment = record_first_segments[record];
+    int64_t offset = record_offsets[record];
+
+    __shared__ AccumulatedScore score_buffers[2][STATE_COUNT];
+    AccumulatedScore* current_scores = score_buffers[0];
+    AccumulatedScore* next_scores = score_buffers[1];
+    if (state < STATE_COUNT) {
+        StateScore initial_score = first_base_score_for_state(
+            state, bases + offset, record_lengths[record],
+            coding_scores[offset], noncoding_scores[offset],
+            transition_probabilities[offset * TRANSITION_COUNT + ENTER_GENE],
+            transition_probabilities[offset * TRANSITION_COUNT + LEAVE_GENE]
+        );
+        current_scores[state] = initial_score == UNREACHABLE_STATE_SCORE
+            ? UNREACHABLE_ACCUMULATED_SCORE
+            : static_cast<AccumulatedScore>(initial_score);
+    }
+    __syncthreads();
+
+    for (int local_segment = 0; local_segment < segment_count; local_segment++) {
+        int segment = first_segment + local_segment;
+        if (state < STATE_COUNT) {
+            segment_entry_scores[segment_state_index(segment, state)] = current_scores[state];
+
+            const StateScore* matrix = transition_matrices
+                + segment_transition_index(segment, state, 0);
+            AccumulatedScore best_score = UNREACHABLE_ACCUMULATED_SCORE;
+#pragma unroll
+            for (int entry_state = 0; entry_state < STATE_COUNT; entry_state++) {
+                AccumulatedScore score_before_segment = current_scores[entry_state];
+                StateScore score_inside_segment = matrix[entry_state];
+                if (score_before_segment == UNREACHABLE_ACCUMULATED_SCORE
+                    || score_inside_segment == UNREACHABLE_STATE_SCORE) {
+                    continue;
+                }
+                AccumulatedScore candidate_score =
+                    score_before_segment
+                    + static_cast<AccumulatedScore>(score_inside_segment);
+                if (candidate_score > best_score) best_score = candidate_score;
+            }
+            next_scores[state] = best_score;
+        }
+        __syncthreads();
+        AccumulatedScore* previous_scores = current_scores;
+        current_scores = next_scores;
+        next_scores = previous_scores;
+    }
+
+    if (state == 0) {
+        int best_end_state = INTERGENIC;
+        AccumulatedScore best_score = current_scores[INTERGENIC];
+        for (int candidate = START_A; candidate <= STOP; candidate++) {
+            if (current_scores[candidate] > best_score) {
+                best_score = current_scores[candidate];
+                best_end_state = candidate;
+            }
+        }
+        if (current_scores[TRAILING_BOUNDARY_INTRON] > best_score) {
+            best_end_state = TRAILING_BOUNDARY_INTRON;
+        }
+        record_end_states[record] = best_end_state;
+    }
+}
+
+__global__ void find_segment_entry_origins(
+    const StateScore* coding_scores, const StateScore* noncoding_scores,
+    const StateScore* transition_probabilities, const int8_t* bases,
+    const int64_t* segment_record_offsets,
+    const int* segment_record_lengths,
+    const int64_t* segment_starts, const int64_t* segment_ends,
+    const AccumulatedScore* segment_entry_scores,
+    int segment_count, int8_t* segment_entry_origins
+) {
+    // A segment starts with the exact 30-state score vector produced by the
+    // preceding segments. Alongside each score, carry only the segment entry
+    // state that produced it. This preserves strict-greater tie order without
+    // storing another per-base backpointer table.
+    int warp_in_block = threadIdx.x >> 5;
+    int segment = blockIdx.x * 4 + warp_in_block;
+    if (segment >= segment_count) return;
+
+    int target_state = threadIdx.x & 31;
+    int64_t record_offset = segment_record_offsets[segment];
+    int record_length = segment_record_lengths[segment];
+    int segment_start = static_cast<int>(segment_starts[segment] - record_offset);
+    int segment_end = static_cast<int>(segment_ends[segment] - record_offset);
+    const StateScore* record_coding_scores = coding_scores + record_offset;
+    const StateScore* record_noncoding_scores = noncoding_scores + record_offset;
+    const StateScore* record_transition_probabilities =
+        transition_probabilities + record_offset * TRANSITION_COUNT;
+    const int8_t* record_bases = bases + record_offset;
+
+    __shared__ AccumulatedScore shared_scores[4][STATE_COUNT];
+    __shared__ int8_t shared_entry_origins[4][STATE_COUNT];
+    AccumulatedScore* current_scores = shared_scores[warp_in_block];
+    int8_t* entry_origins = shared_entry_origins[warp_in_block];
+    if (target_state < STATE_COUNT) {
+        current_scores[target_state] =
+            segment_entry_scores[segment_state_index(segment, target_state)];
+        entry_origins[target_state] = static_cast<int8_t>(target_state);
+    }
+    __syncwarp();
+
+    for (int position = segment_start + 1; position <= segment_end; position++) {
+        PositionSignals signals = compute_position_signals(
+            record_coding_scores, record_noncoding_scores,
+            record_transition_probabilities, record_bases, record_length, position
+        );
+        TransitionChoice<AccumulatedScore> best = {UNREACHABLE_ACCUMULATED_SCORE, 0};
+        int8_t best_entry_origin = 0;
+        if (target_state < STATE_COUNT) {
+            best = best_transition_to_state(target_state, current_scores, signals);
+            if (!is_unreachable(best.score)) {
+                best_entry_origin = entry_origins[best.source_state];
+            }
+        }
+
+        if (target_state < STATE_COUNT) {
+            current_scores[target_state] = best.score;
+            entry_origins[target_state] = best_entry_origin;
+        }
+        __syncwarp();
+    }
+
+    if (target_state < STATE_COUNT) {
+        segment_entry_origins[segment_state_index(segment, target_state)] =
+            entry_origins[target_state];
+    }
+}
+
+__global__ void trace_selected_segment_boundaries(
+    const int* record_first_segments, const int* record_segment_counts,
+    const int* record_end_states,
+    const int8_t* segment_entry_origins,
+    int record_count, int8_t* segment_entry_states, int* segment_exit_states
+) {
+    // Follow the exact segment-origin table from the selected record end state.
+    int record = blockIdx.x * blockDim.x + threadIdx.x;
+    if (record >= record_count) return;
+
+    int first_segment = record_first_segments[record];
+    int segment_count = record_segment_counts[record];
+    int current_exit_state = record_end_states[record];
+    for (int local_segment = segment_count - 1; local_segment >= 0; local_segment--) {
+        int segment = first_segment + local_segment;
+        segment_exit_states[segment] = current_exit_state;
+        current_exit_state = static_cast<int>(
+            segment_entry_origins[
+                segment_state_index(segment, current_exit_state)
+            ]
+        );
+        segment_entry_states[segment] = static_cast<int8_t>(current_exit_state);
+    }
+}
+
+// Reused device buffers grow as needed and remain allocated between batches.
+static __half* d_probabilities = nullptr;
+static StateScore* d_coding_scores = nullptr;
+static StateScore* d_noncoding_scores = nullptr;
+static StateScore* d_transition_probabilities = nullptr;
+static StateScore* d_transition_matrices = nullptr;
+static AccumulatedScore* d_segment_entry_scores = nullptr;
+static int8_t* d_bases = nullptr;
+static int8_t* d_backpointers = nullptr;
+static int8_t* d_paths = nullptr;
+static int8_t* d_segment_entry_origins = nullptr;
+static int8_t* d_segment_entry_states = nullptr;
+static int64_t* d_record_offsets = nullptr;
+static int64_t* d_segment_starts = nullptr;
+static int64_t* d_segment_ends = nullptr;
+static int64_t* d_segment_record_offsets = nullptr;
+static int* d_record_lengths = nullptr;
+static int* d_segment_exit_states = nullptr;
+static int* d_segment_record_lengths = nullptr;
+static int* d_record_end_states = nullptr;
+static int* d_record_first_segments = nullptr;
+static int* d_record_segment_counts = nullptr;
+
+static int64_t allocated_base_capacity = 0;
+static int64_t allocated_segment_capacity = 0;
+static int64_t allocated_record_capacity = 0;
+static int64_t allocated_backpointer_capacity = 0;
+
+template <typename T>
+static void reallocate_device_buffer(T*& buffer, size_t element_count) {
+    if (buffer) cudaFree(buffer);
+    cudaMalloc(reinterpret_cast<void**>(&buffer), element_count * sizeof(T));
+}
+
+template <typename T>
+static void copy_to_device(T* destination, const T* source, size_t element_count) {
+    cudaMemcpy(destination, source, element_count * sizeof(T), cudaMemcpyHostToDevice);
+}
+
+extern "C" void decode_batch_on_gpu(
+    const unsigned short* probability_bits, const int8_t* input_bases,
+    const StateScore* input_transition_probabilities,
+    const int64_t* input_record_offsets, int record_count,
+    int requested_segment_length,
+    int64_t requested_reconstruction_batch_size,
+    int8_t* paths
+) {
+    const __half* probabilities = reinterpret_cast<const __half*>(probability_bits);
+    if (record_count <= 0) return;
+
+    // One step can add at most one transition log-odds value, so the largest
+    // magnitude widens the per-step bound that sizes the int32 segments.
+    int64_t largest_transition_log_odds = 0;
+    int64_t total_bases = input_record_offsets[record_count];
+    for (int64_t value = 0; value < total_bases * TRANSITION_COUNT; value++) {
+        int64_t magnitude = std::abs(static_cast<int64_t>(input_transition_probabilities[value]));
+        if (magnitude > largest_transition_log_odds) largest_transition_log_odds = magnitude;
+    }
+    int segment_length = requested_segment_length < 1 ? 2000 : requested_segment_length;
+    int64_t maximum_batch_step_score =
+        maximum_absolute_step_score + largest_transition_log_odds;
+    if (maximum_batch_step_score >= RAW_SEGMENT_SCORE_LIMIT) {
+        throw std::invalid_argument(
+            "combined decoder and transition probabilities exceed the int32 step range"
+        );
+    }
+    int maximum_safe_segment_length = static_cast<int>(
+        RAW_SEGMENT_SCORE_LIMIT / maximum_batch_step_score
+    );
+    if (segment_length > maximum_safe_segment_length) {
+        segment_length = maximum_safe_segment_length;
+    }
+    // This is a base count for the final path-reconstruction batch, not a
+    // segment count and not a byte count. Zero selects automatic sizing.
+    int64_t reconstruction_batch_size = requested_reconstruction_batch_size;
+    if (reconstruction_batch_size == 0) {
+        size_t free_bytes = 0, total_bytes = 0;
+        cudaMemGetInfo(&free_bytes, &total_bytes);
+        reconstruction_batch_size = static_cast<int64_t>(
+            free_bytes / 8 / BACKPOINTER_STRIDE
+        );
+    }
+    if (reconstruction_batch_size < static_cast<int64_t>(segment_length) * 2) {
+        reconstruction_batch_size = static_cast<int64_t>(segment_length) * 2;
+    }
+    int64_t required_backpointer_capacity = (
+        total_bases < reconstruction_batch_size
+            ? total_bases
+            : reconstruction_batch_size
+    );
+
+    std::vector<int> record_lengths(record_count);
+    std::vector<int> record_first_segments(record_count), record_segment_counts(record_count);
+    int segment_count = 0;
+    // All segments use the same algorithm and are distributed as evenly as
+    // possible; there is no special segment-zero decoder.
+    for (int record = 0; record < record_count; record++) {
+        int64_t length = input_record_offsets[record + 1] - input_record_offsets[record];
+        record_lengths[record] = static_cast<int>(length);
+        int count = static_cast<int>((length - 1 + segment_length - 1) / segment_length);
+        if (count < 1) count = 1;
+        record_first_segments[record] = segment_count;
+        record_segment_counts[record] = count;
+        segment_count += count;
+    }
+
+    // Persistent device buffers grow on demand and are reused by later batches.
+    if (total_bases > allocated_base_capacity) {
+        reallocate_device_buffer(d_probabilities, total_bases);
+        reallocate_device_buffer(d_coding_scores, total_bases);
+        reallocate_device_buffer(d_noncoding_scores, total_bases);
+        reallocate_device_buffer(
+            d_transition_probabilities,
+            static_cast<size_t>(total_bases) * TRANSITION_COUNT
+        );
+        reallocate_device_buffer(d_bases, total_bases);
+        reallocate_device_buffer(d_paths, total_bases);
+        allocated_base_capacity = total_bases;
+    }
+    if (required_backpointer_capacity > allocated_backpointer_capacity) {
+        reallocate_device_buffer(
+            d_backpointers, static_cast<size_t>(required_backpointer_capacity) * BACKPOINTER_STRIDE
+        );
+        allocated_backpointer_capacity = required_backpointer_capacity;
+    }
+    if (record_count > allocated_record_capacity) {
+        reallocate_device_buffer(d_record_offsets, record_count + 1);
+        reallocate_device_buffer(d_record_lengths, record_count);
+        reallocate_device_buffer(d_record_end_states, record_count);
+        reallocate_device_buffer(d_record_first_segments, record_count);
+        reallocate_device_buffer(d_record_segment_counts, record_count);
+        allocated_record_capacity = record_count;
+    }
+    if (segment_count > allocated_segment_capacity) {
+        reallocate_device_buffer(d_segment_starts, segment_count);
+        reallocate_device_buffer(d_segment_ends, segment_count);
+        reallocate_device_buffer(d_segment_record_offsets, segment_count);
+        reallocate_device_buffer(d_segment_record_lengths, segment_count);
+        reallocate_device_buffer(d_segment_exit_states, segment_count);
+        reallocate_device_buffer(
+            d_transition_matrices, static_cast<size_t>(segment_count) * STATE_COUNT * STATE_COUNT
+        );
+        reallocate_device_buffer(
+            d_segment_entry_scores, static_cast<size_t>(segment_count) * STATE_COUNT
+        );
+        reallocate_device_buffer(
+            d_segment_entry_origins, static_cast<size_t>(segment_count) * STATE_COUNT
+        );
+        reallocate_device_buffer(d_segment_entry_states, segment_count);
+        allocated_segment_capacity = segment_count;
+    }
+
+    // Flatten each record's segments while retaining the record-relative data
+    // needed by matrix construction and backpointer recomputation.
+    std::vector<int64_t> segment_starts(segment_count), segment_ends(segment_count);
+    std::vector<int64_t> segment_record_offsets(segment_count);
+    std::vector<int> segment_record_lengths(segment_count);
+    for (int record = 0; record < record_count; record++) {
+        int64_t offset = input_record_offsets[record];
+        int length = record_lengths[record];
+        int first_segment = record_first_segments[record];
+        int count = record_segment_counts[record];
+        auto boundary = [&](int boundary_index) {
+            return static_cast<int64_t>(
+                static_cast<int64_t>(boundary_index) * (length - 1) / count
+            );
+        };
+
+        for (int local_segment = 0; local_segment < count; local_segment++) {
+            int segment = first_segment + local_segment;
+            segment_starts[segment] = offset + boundary(local_segment);
+            segment_ends[segment] = offset + boundary(local_segment + 1);
+            segment_record_offsets[segment] = offset;
+            segment_record_lengths[segment] = length;
+        }
+    }
+
+    // Upload packed per-base data and the compact segment layout.
+    cudaMemcpy(d_probabilities, probability_bits, total_bases * 2, cudaMemcpyHostToDevice);
+    copy_to_device(
+        d_transition_probabilities, input_transition_probabilities,
+        static_cast<size_t>(total_bases) * TRANSITION_COUNT
+    );
+    copy_to_device(d_bases, input_bases, total_bases);
+    copy_to_device(d_record_offsets, input_record_offsets, record_count + 1);
+    copy_to_device(d_record_lengths, record_lengths.data(), record_count);
+    copy_to_device(d_record_first_segments, record_first_segments.data(), record_count);
+    copy_to_device(d_record_segment_counts, record_segment_counts.data(), record_count);
+    copy_to_device(d_segment_starts, segment_starts.data(), segment_count);
+    copy_to_device(d_segment_ends, segment_ends.data(), segment_count);
+    copy_to_device(d_segment_record_offsets, segment_record_offsets.data(), segment_count);
+    copy_to_device(d_segment_record_lengths, segment_record_lengths.data(), segment_count);
+
+    // Build every segment transition matrix, then combine matrices by record.
+    int threads = 256;
+    int probability_blocks = static_cast<int>((total_bases + threads - 1) / threads);
+    convert_probabilities_to_log_scores<<<probability_blocks, threads>>>(
+        d_probabilities, d_coding_scores, d_noncoding_scores,
+        static_cast<int>(total_bases)
+    );
+    int64_t matrix_threads = static_cast<int64_t>(segment_count) * STATE_COUNT;
+    int matrix_blocks = static_cast<int>((matrix_threads + 127) / 128);
+    build_segment_transition_matrices<<<matrix_blocks, 128>>>(
+        d_coding_scores, d_noncoding_scores, d_transition_probabilities,
+        d_bases, d_segment_record_offsets,
+        d_segment_record_lengths, d_segment_starts, d_segment_ends,
+        segment_count, d_transition_matrices
+    );
+    combine_segment_scores<<<record_count, 32>>>(
+        d_coding_scores, d_noncoding_scores, d_transition_probabilities,
+        d_bases, d_record_offsets, d_record_lengths,
+        d_transition_matrices, d_record_first_segments,
+        d_record_segment_counts, record_count,
+        d_segment_entry_scores, d_record_end_states
+    );
+    find_segment_entry_origins<<<(segment_count + 3) / 4, 128>>>(
+        d_coding_scores, d_noncoding_scores, d_transition_probabilities, d_bases,
+        d_segment_record_offsets, d_segment_record_lengths,
+        d_segment_starts, d_segment_ends,
+        d_segment_entry_scores, segment_count,
+        d_segment_entry_origins
+    );
+
+    trace_selected_segment_boundaries<<<(record_count + 127) / 128, 128>>>(
+        d_record_first_segments, d_record_segment_counts,
+        d_record_end_states, d_segment_entry_origins, record_count,
+        d_segment_entry_states, d_segment_exit_states
+    );
+
+    // Selected-path reconstruction materializes backpointers in bounded
+    // base-pair windows. Each selected segment restarts at score 0 and uses
+    // raw int32 scores without a clamp.
+    std::vector<int> window_first_segments;
+    {
+        int segment = 0;
+        while (segment < segment_count) {
+            int64_t window_start_bp = segment_starts[segment];
+            int end = segment + 1;
+            while (end < segment_count
+                   && segment_ends[end] - window_start_bp <= required_backpointer_capacity) {
+                end++;
+            }
+            window_first_segments.push_back(segment);
+            segment = end;
+        }
+        window_first_segments.push_back(segment_count);
+    }
+    int window_count = static_cast<int>(window_first_segments.size()) - 1;
+
+    auto recompute_selected_backpointers = [&](int start, int end, int64_t window_start_bp) {
+        recompute_selected_segment_backpointers<<<(end - start + 3) / 4, 128>>>(
+            d_coding_scores, d_noncoding_scores, d_transition_probabilities,
+            d_bases, d_segment_record_offsets + start,
+            d_segment_record_lengths + start, d_segment_starts + start,
+            d_segment_ends + start,
+            d_segment_entry_states + start,
+            end - start, window_start_bp, d_backpointers
+        );
+    };
+    auto write_selected_paths = [&](int start, int end, int64_t window_start_bp) {
+        write_selected_segment_paths<<<(end - start + 127) / 128, 128>>>(
+            d_segment_starts + start, d_segment_ends + start,
+            d_segment_exit_states + start, end - start,
+            window_start_bp, d_backpointers, d_paths
+        );
+    };
+
+    for (int window = 0; window < window_count; window++) {
+        int start = window_first_segments[window];
+        int end = window_first_segments[window + 1];
+        int64_t window_start_bp = segment_starts[start];
+        recompute_selected_backpointers(start, end, window_start_bp);
+        write_selected_paths(start, end, window_start_bp);
+    }
+
+    cudaMemcpy(paths, d_paths, total_bases, cudaMemcpyDeviceToHost);
+    cudaDeviceSynchronize();
+}
+
+#ifdef TORCH_EXTENSION_NAME
+#include <torch/extension.h>
+#include <pybind11/numpy.h>
+namespace py = pybind11;
+
+static py::array_t<int8_t> decode_packed_paths(
+    py::array_t<unsigned short> probability_array, py::array_t<int8_t> encoded_bases,
+    py::array_t<StateScore> transition_probabilities,
+    py::array_t<int64_t> record_offsets_array,
+    int segment_length, int64_t reconstruction_batch_size
+) {
+    if (transition_probabilities.size() != encoded_bases.size() * TRANSITION_COUNT) {
+        throw std::invalid_argument(
+            "transition probabilities must contain four values per encoded base"
+        );
+    }
+    int record_count = static_cast<int>(record_offsets_array.size()) - 1;
+    auto state_paths = py::array_t<int8_t>(encoded_bases.size());
+    const unsigned short* probability_bits = probability_array.data();
+    const int8_t* base_data = encoded_bases.data();
+    const StateScore* transition_probability_data = transition_probabilities.data();
+    const int64_t* record_offsets = record_offsets_array.data();
+    int8_t* output_paths = state_paths.mutable_data();
+    {
+        // Let the Python reader prepare the next input chunk during GPU work.
+        py::gil_scoped_release release;
+        decode_batch_on_gpu(
+            probability_bits, base_data, transition_probability_data, record_offsets,
+            record_count, segment_length, reconstruction_batch_size, output_paths
+        );
+    }
+    return state_paths;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("init", &initialize_gpu_decoder, "Set decoder stop codons, intron priors, and limits");
+    m.def("decode", &decode_packed_paths, "Decode a batch with the 30-state CDS decoder");
+}
+#endif

annotation_pipeline/weights/confidence_network.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eadd22818e30aee6e8e8d779b3b25302bd503ce42dbb82b3f9df9620c9089928
+size 63460

annotation_pipeline/weights/transition_network.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e4d253ab067a97eaeee690c1292000d95ea649326f6bad66f4570f16fb138f9
+size 49628

genbank-inference-check.json

@@ -1,33 +0,0 @@
-{
-  "status": "passed",
-  "model_revision": "d429145e77373827ac93cc03ff04dcbdfab5723f",
-  "gpu": "NVIDIA H100 80GB HBM3",
-  "torch": "2.7.1+cu128",
-  "dtype": "bfloat16",
-  "source_accession": "NC_001133.9",
-  "source_length": 230218,
-  "source_windows": 4,
-  "records_per_format": 3,
-  "formats": [
-    "gzipped GenBank",
-    "gzipped FASTA"
-  ],
-  "checks": [
-    "published script and model downloaded",
-    "real NCBI chromosome inference",
-    "multiple records",
-    "all original coordinates retained",
-    "independent overlap-average comparison",
-    "GenBank and FASTA predictions identical",
-    "ambiguous 6-mers excluded",
-    "incomplete final 6-mer excluded",
-    "single-base record handled",
-    "thresholded strand and combined calls",
-    "overwrite refusal"
-  ],
-  "scope": "Functional inference check; not an accuracy benchmark",
-  "date": "2026-10-02",
-  "transformers": "4.56.0",
-  "biopython": "1.85",
-  "slurm_job_id": "100266"
-}

infer_genbank.py

@@ -1,123 +0,0 @@
-"""Predict per-base CDS occupancy from a GenBank/GBFF (or FASTA) file."""
-import argparse
-import gzip
-import json
-from pathlib import Path
-import re
-
-from Bio import SeqIO
-import numpy as np
-import torch
-from transformers import AutoModelForTokenClassification, AutoTokenizer
-
-MODEL_ID = 'HuggingFaceBio/Carbon-A-1.2B'
-WINDOW = 98304
-STRIDE = 49152
-
-
-def safe_name(value):
-    return re.sub(r'[^A-Za-z0-9_.-]+', '_', value).strip('._')[:100] or 'record'
-
-
-def predict_record(sequence, model, tokenizer, device, threshold):
-    """Average probabilities across overlapping windows in record coordinates."""
-    sequence = sequence.upper()
-    length = len(sequence)
-    if not length:
-        raise ValueError('Cannot predict an empty sequence')
-    sums = np.zeros((2, length), dtype=np.float32)
-    counts = np.zeros(length, dtype=np.uint32)
-    windows = 0
-    with torch.inference_mode():
-        for start in range(0, length, STRIDE):
-            stop = min(start + WINDOW, length)
-            fragment = sequence[start:stop]
-            # Complete the final 6-mer with N; its positions will be unscored.
-            fragment += 'N' * (-len(fragment) % 6)
-            encoded = tokenizer(fragment, add_special_tokens=False, return_tensors='pt',
-                                return_attention_mask=True)
-            ids = encoded['input_ids']
-            n_tokens = len(fragment) // 6
-            if ids.shape != (1, n_tokens):
-                raise ValueError('Tokenizer did not produce one token per six bases')
-            # Match packed-data inference: OOV tokens remain visible context;
-            # only their output positions are excluded, not the attention mask.
-            mask = torch.ones_like(ids)
-            logits = model(input_ids=ids.to(device), attention_mask=mask.to(device)).logits.float()
-            if logits.shape != (1, 2 * n_tokens * 6, 2) or not torch.isfinite(logits).all():
-                raise ValueError('Unexpected shape or nonfinite model predictions')
-            probs = logits.softmax(-1)[0, :, 1].cpu().numpy().reshape(2, -1)
-            width = stop - start
-            valid = np.repeat(ids[0].numpy() >= 32, 6)[:width]
-            sums[:, start:stop] += np.where(valid[None, :], probs[:, :width], 0)
-            counts[start:stop] += valid.astype(np.uint32)
-            windows += 1
-            if stop == length:
-                break
-    valid = counts > 0
-    np.divide(sums, counts[None, :], out=sums, where=valid[None, :])
-    sums[:, ~valid] = np.nan
-    calls = (sums >= threshold).astype(np.int8)
-    calls[:, ~valid] = -1
-    any_call = np.any(calls == 1, axis=0).astype(np.int8)
-    any_call[~valid] = -1
-    return {'cds_probability_plus': sums[0], 'cds_probability_minus': sums[1],
-            'cds_label_plus': calls[0], 'cds_label_minus': calls[1],
-            'cds_label_any': any_call, 'valid': valid, 'window_count': counts}, windows
-
-
-def main():
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument('input', type=Path, help='GenBank/GBFF file; .gz is supported.')
-    parser.add_argument('--format', choices=['genbank', 'fasta'], default='genbank')
-    parser.add_argument('--model', default=MODEL_ID)
-    parser.add_argument('--revision', default='main', help='Use a model commit for repeatable runs.')
-    parser.add_argument('--device', choices=['cuda', 'cpu'], default='cuda')
-    parser.add_argument('--threshold', type=float, default=0.5)
-    parser.add_argument('--out-dir', type=Path, default=Path('genbank-predictions'))
-    args = parser.parse_args()
-    if not 0 <= args.threshold <= 1:
-        parser.error('Threshold must be between 0 and 1.')
-    if not args.input.is_file():
-        parser.error('Input file does not exist.')
-    if args.device == 'cuda' and not torch.cuda.is_available():
-        parser.error('CUDA is unavailable; use a GPU or explicitly select --device cpu (slow).')
-    if args.out_dir.exists() and (not args.out_dir.is_dir() or any(args.out_dir.iterdir())):
-        parser.error('Choose an empty output directory to avoid overwriting previous predictions.')
-    dtype = torch.float32
-    if args.device == 'cuda':
-        dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
-    tokenizer = AutoTokenizer.from_pretrained(args.model, revision=args.revision, trust_remote_code=True)
-    model = AutoModelForTokenClassification.from_pretrained(
-        args.model, revision=args.revision, trust_remote_code=True,
-        dtype=dtype, attn_implementation='sdpa',
-    ).to(args.device).eval()
-    if (model.config.k, model.config.num_prediction_heads, model.config.num_labels) != (6, 2, 2):
-        raise ValueError('Expected a 6-mer, two-strand binary CDS model.')
-    args.out_dir.mkdir(parents=True, exist_ok=True)
-    summary = {'model':args.model, 'revision_requested':args.revision,
-               'revision_resolved':getattr(model.config, '_commit_hash', None),
-               'input_file':str(args.input), 'input_format':args.format,
-               'dtype':str(dtype), 'threshold':args.threshold,
-               'window_bp':WINDOW, 'stride_bp':STRIDE,
-               'coordinates':'Array index i corresponds to record base i+1 on both strands.',
-               'records':[]}
-    opener = gzip.open if args.input.suffix.lower() == '.gz' else open
-    with opener(args.input, 'rt') as handle:
-        for index, record in enumerate(SeqIO.parse(handle, args.format)):
-            if not len(record):
-                raise ValueError(f'Empty sequence in record {record.id}')
-            arrays, windows = predict_record(str(record.seq), model, tokenizer, args.device, args.threshold)
-            filename = f'{index:04d}_{safe_name(record.id)}_{safe_name(record.name)}.npz'
-            np.savez_compressed(args.out_dir / filename, **arrays)
-            summary['records'].append({'accession':record.id, 'name':record.name,
-                                       'length':len(record), 'windows':windows,
-                                       'valid_bases':int(arrays['valid'].sum()), 'file':filename})
-            print(f'Saved {filename}: {len(record):,} bases, {windows} windows', flush=True)
-    if not summary['records']:
-        raise ValueError('No sequence records found; check the input format.')
-    (args.out_dir / 'summary.json').write_text(json.dumps(summary, indent=2)+'\n')
-
-
-if __name__ == '__main__':
-    main()

infer_packed.py

@@ -1,116 +0,0 @@
-"""Run Carbon-A on local packed HDF5 rows, or download a small example file."""
-import argparse
-import json
-from pathlib import Path
-import time
-
-import numpy as np
-import requests
-import torch
-from huggingface_hub import get_token
-from torch.utils.data import DataLoader, Subset
-from transformers import AutoModelForTokenClassification
-
-from load_data import PackedAnnotationDataset
-
-MODEL_ID = 'HuggingFaceBio/Carbon-A-1.2B'
-EXAMPLE_URL = ('https://huggingface.co/buckets/HuggingFaceBio/Carbon-A-training-data/'
-               'resolve/test/fungi1/GCF_000002945.2.h5')
-
-
-def download_example(path):
-    """Stream a 29 MB evaluation file using the saved Hugging Face login."""
-    path.parent.mkdir(parents=True, exist_ok=True)
-    if not path.exists():
-        token = get_token()
-        headers = {'Authorization': f'Bearer {token}'} if token else {}
-        partial = path.with_suffix('.h5.partial')
-        with requests.get(EXAMPLE_URL, headers=headers, stream=True, timeout=(30, 120)) as response:
-            response.raise_for_status()
-            with partial.open('wb') as f:
-                for block in response.iter_content(1024 * 1024):
-                    f.write(block)
-        partial.replace(path)
-    return path
-
-
-def main():
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument('--model', default=MODEL_ID)
-    parser.add_argument('--revision', default='main', help='Pin a model commit for repeatable runs.')
-    parser.add_argument('--h5', type=Path, help='Local packed HDF5. If omitted, download the example evaluation file.')
-    parser.add_argument('--cache-dir', type=Path, default=Path('data'))
-    parser.add_argument('--start-row', type=int, default=0)
-    parser.add_argument('--num-rows', type=int, default=1)
-    parser.add_argument('--batch-size', type=int, default=1)
-    parser.add_argument('--device', choices=['cuda', 'cpu'], default='cuda')
-    parser.add_argument('--dtype', choices=['auto', 'float32', 'bfloat16', 'float16'], default='auto')
-    parser.add_argument('--threshold', type=float, default=0.5)
-    parser.add_argument('--out-dir', type=Path, default=Path('predictions'))
-    args = parser.parse_args()
-    if args.start_row < 0 or args.num_rows < 1 or args.batch_size < 1:
-        parser.error('Use a nonnegative start row and positive row/batch counts.')
-    if not 0 <= args.threshold <= 1:
-        parser.error('Threshold must be between 0 and 1.')
-    if args.device == 'cuda' and not torch.cuda.is_available():
-        parser.error('CUDA is unavailable. Run on a GPU, or explicitly select --device cpu (slow).')
-    dtype_name = args.dtype
-    if dtype_name == 'auto':
-        dtype_name = ('bfloat16' if torch.cuda.is_bf16_supported() else 'float16') if args.device == 'cuda' else 'float32'
-    dtype = getattr(torch, dtype_name)
-    path = args.h5 or download_example(args.cache_dir / 'GCF_000002945.2.h5')
-    dataset = PackedAnnotationDataset([path])
-    stop = args.start_row + args.num_rows
-    if stop > len(dataset):
-        parser.error(f'Requested rows end at {stop}; file contains {len(dataset)} rows.')
-    args.out_dir.mkdir(parents=True, exist_ok=True)
-    targets = [args.out_dir / f'{path.stem}.row-{i:06d}.npz' for i in range(args.start_row, stop)]
-    report_path = args.out_dir / f'{path.stem}.rows-{args.start_row}-{stop}.json'
-    if any(p.exists() for p in targets + [report_path]):
-        parser.error('Output already exists; choose a new output directory or row range.')
-    model = AutoModelForTokenClassification.from_pretrained(
-        args.model, revision=args.revision, trust_remote_code=True,
-        dtype=dtype, attn_implementation='sdpa',
-    ).to(args.device).eval()
-    if (model.config.k, model.config.num_prediction_heads, model.config.num_labels) != (6, 2, 2):
-        raise ValueError('Expected a 6-mer, two-strand binary CDS model.')
-    loader = DataLoader(Subset(dataset, range(args.start_row, stop)), batch_size=args.batch_size)
-    reports = []
-    cursor = 0
-    started = time.monotonic()
-    with torch.inference_mode():
-        for batch in loader:
-            logits = model(input_ids=batch['input_ids'].to(args.device),
-                           attention_mask=batch['attention_mask'].to(args.device)).logits.float()
-            if logits.shape[1:] != (196608, 2) or not torch.isfinite(logits).all():
-                raise ValueError(f'Unexpected or nonfinite logits: {tuple(logits.shape)}')
-            probs = logits.softmax(-1)[..., 1].cpu().numpy().reshape(-1, 2, 98304)
-            references = batch['labels'].numpy().reshape(-1, 2, 98304)
-            for prob, reference in zip(probs, references):
-                valid = reference != -100
-                calls = (prob >= args.threshold).astype(np.int8)
-                calls[~valid] = -1
-                any_valid = valid.all(axis=0)
-                any_call = np.any(calls == 1, axis=0).astype(np.int8)
-                any_call[~any_valid] = -1
-                output = targets[cursor]
-                np.savez_compressed(output, cds_probability_plus=prob[0], cds_probability_minus=prob[1],
-                                    cds_label_plus=calls[0], cds_label_minus=calls[1], cds_label_any=any_call,
-                                    valid_plus=valid[0], valid_minus=valid[1],
-                                    reference_plus=reference[0], reference_minus=reference[1])
-                reports.append({'row':args.start_row+cursor, 'file':output.name,
-                                'bases_per_strand':98304, 'valid_bases_per_strand':valid.sum(axis=1).tolist(),
-                                'cds_calls_per_strand':(calls == 1).sum(axis=1).tolist()})
-                cursor += 1
-                print(f'Saved {output}', flush=True)
-    report = {'model':args.model, 'revision_requested':args.revision,
-              'revision_resolved':getattr(model.config, '_commit_hash', None),
-              'input_file':str(path), 'device':args.device, 'dtype':dtype_name,
-              'threshold':args.threshold, 'elapsed_seconds':time.monotonic()-started,
-              'rows':reports}
-    report_path.write_text(json.dumps(report, indent=2)+'\n')
-    print(json.dumps(report, indent=2), flush=True)
-
-
-if __name__ == '__main__':
-    main()

inference-check.json

@@ -1,27 +0,0 @@
-{
-  "status": "passed",
-  "gpu": "NVIDIA H100 80GB HBM3",
-  "torch": "2.7.1+cu128",
-  "full_length_rows": 2,
-  "batch_size": 2,
-  "bases_per_strand": 98304,
-  "checks": [
-    "published model and scripts downloaded",
-    "bucket example downloaded by inference script",
-    "finite logits and probabilities",
-    "two strand shapes",
-    "threshold calls and validity masks",
-    "combined strand calls",
-    "tokenizer IDs"
-  ],
-  "tested_model_revision": "d5a62f1bbf483426314525fcff8695847f9fcdc4",
-  "transformers": "4.56.0",
-  "dtype": "bfloat16",
-  "data_bucket": "HuggingFaceBio/Carbon-A-training-data",
-  "data_path": "test/fungi1/GCF_000002945.2.h5",
-  "rows": [
-    0,
-    1
-  ],
-  "scope": "Functional inference check, not an accuracy benchmark"
-}

load_data.py

@@ -1,52 +0,0 @@
-"""Minimal loader for the released 16,384-token GENERanno annotation HDF5s."""
-from bisect import bisect_right
-from pathlib import Path
-
-import h5py
-import numpy as np
-import torch
-from torch.utils.data import Dataset
-
-
-class PackedAnnotationDataset(Dataset):
-    """Read local downloaded shards; one row = 98,304 bp with two label tracks.
-
-    No augmentation, retokenization, or context truncation is applied here.
-    Supply special IDs from the bundled tokenizer to reproduce training masking.
-    """
-
-    def __init__(self, paths, ignored_token_ids=(0, 1, 2, 3, 4)):
-        self.paths = sorted(Path(p) for p in paths)
-        if not self.paths:
-            raise ValueError("Provide at least one downloaded .h5 file")
-        self.ignored_token_ids = np.asarray(ignored_token_ids, dtype=np.int64)
-        self.ends = []
-        total = 0
-        for path in self.paths:
-            with h5py.File(path, "r") as f:
-                n, tokens = f['input_ids'].shape
-                if tokens != 16384 or any(f[k].shape != (n, tokens * 6) for k in ('label_plus', 'label_minus')):
-                    raise ValueError(f"Expected aligned 16k annotation rows in {path}")
-                total += n
-                self.ends.append(total)
-
-    def __len__(self):
-        return self.ends[-1]
-
-    def __getitem__(self, index):
-        if not 0 <= index < len(self):
-            raise IndexError(index)
-        file_index = bisect_right(self.ends, index)
-        row = index - (self.ends[file_index - 1] if file_index else 0)
-        # Open per read: no shared HDF5 handle between DataLoader workers.
-        with h5py.File(self.paths[file_index], "r") as f:
-            ids = np.asarray(f['input_ids'][row], dtype=np.int64)
-            plus = (f['label_plus'][row] != 0).astype(np.int64)
-            minus = (f['label_minus'][row] != 0).astype(np.int64)
-        labels = np.concatenate([plus, minus])
-        ignored_bases = np.repeat(np.isin(ids, self.ignored_token_ids), 6)
-        labels[np.concatenate([ignored_bases, ignored_bases])] = -100
-        return {'input_ids': torch.from_numpy(ids),
-                # Fixed-length pretokenized rows; matches the original collator.
-                'attention_mask': torch.ones(len(ids), dtype=torch.long),
-                'labels': torch.from_numpy(labels)}

provenance.json

@@ -37,12 +37,7 @@
   "additional_files": [
     "README.md",
     "requirements.txt",
-    "load_data.py",
-    "infer_packed.py",
-    "infer_genbank.py",
-    "inference-check.json",
-    "genbank-inference-check.json",
-    "INFERENCE.md"
+    "annotation_pipeline/"
   ],
   "modified_source_files": {
     "modeling_generanno.py": {

requirements.txt

@@ -1,8 +1,20 @@
 torch>=2.7,<3
-transformers>=4.56,<6
+# Below 5: Transformers 5 does not load this model with flash_attention_2, the
+# annotation pipeline's default attention implementation.
+transformers>=4.56,<5
 huggingface_hub>=0.34
 safetensors>=0.4.3
 numpy>=1.24
-h5py>=3.10
-requests>=2.31
 biopython>=1.83,<2
+
+# Annotation pipeline (annotation_pipeline/, Linux only)
+pandas
+pyarrow
+tqdm
+numba
+ninja
+# Its default attention implementation, flash_attention_2, also needs FlashAttention:
+# pip install --no-build-isolation flash-attn
+# (--attn-implementation sdpa runs without it.) The native Viterbi backends are
+# compiled on first use: the CPU backend needs a C++17 compiler with OpenMP and
+# the GPU backend a CUDA toolkit.