GenerTeam's picture
Support CPU inference
56c0e2f verified
|
Raw History Blame Contribute Delete
12.7 kB

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:

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:

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:

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 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 --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 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; 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:

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:

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:

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.