# 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 a model-derived confidence score
   for the decoded gene structure (`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. Without a
GPU, a CPU build of PyTorch suffices and FlashAttention is not needed.

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. Without a GPU, or with
`--gpu-count 0`, it runs on the CPU, in float32 with SDPA. On the CPU, inference
and post-processing use every core available to them; `OMP_NUM_THREADS` sets
the number of threads. `--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 score for the decoded gene
structure, or `confidence=NA` for genes truncated by a record edge. Higher
scores indicate greater model confidence. 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 `--tcn-device DEVICE`
picks the torch device for the two networks, both built on a temporal
convolutional network (TCN); the default is `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. The check groups
the complete genes by codon: the genes ending in each stop codon of the table,
and the genes containing in frame each of TAA, TAG and TGA that the table reads
as an amino acid. When, in any group, significantly fewer than half reach
confidence 0.1 (one-sided binomial test against one half, p < 0.01), it warns
that the table probably does not match. The warning names each such group and
goes to stderr, to `WARNINGS.txt` in the parent of the output directory, and
into the output directory's name, as described above.

## 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`.