Match genome-scale inference windowing and batching

#3
by kashif HF Staff - opened
Hugging Face Biology Research org
•
edited 2 days ago

Makes infer_genbank.py reproduce the inference used for Carbon-A's genome-scale predictions.

  • Windows: 16,384 tokens with a 1,024-token overlap (--overlap-length) instead of 50%. The last window is right-aligned to the record end, so no N bases are added and no tail bases are lost. Records shorter than one window are run aligned to their start and to their end, and records under 6 bp get P(CDS) = 0. Compute drops from about 2× to about 1.07× per base.
  • Batching: records longest first, in groups of 4 (--dispatch-chunk-size); full-length and padded windows in separate batches of 8 (--batch-size). Softmax in the model dtype, with P(CDS) summed and divided by coverage. In BF16, batch composition changes results slightly, so this is needed for exact reproduction.
  • Ambiguous 6-mers are predicted instead of being set to NaN. An ambiguous mask replaces valid.
  • --attn-implementation flash_attention_2 option, matching the attention kernel used for the genome-scale predictions.
  • README: the Python example uses predict_record instead of N-padding. README and INFERENCE.md describe the windowing, batching, outputs and exactness.

Checked on an RTX 3090 Ti against the reference inference, on the same FASTA files (5 records, and 12 records from 5 bp to 230 kb with Ns):

  • FA2 + BF16, transformers 4.56.0, flash-attn 2.8.3, default batch settings: bit-identical for every record.
  • SDPA on both sides: bit-identical for every record.
  • Script default (SDPA, transformers 5.19) vs the FA2 reference: up to 0.11 on single bases, with ≥ 99.8% of 0.5 calls unchanged.
  • GenBank and FASTA input give identical output.

🤖 Generated with Claude Code

kashif changed pull request title from Match pipeline windowing to Match genome-scale inference windowing and batching
GenerTeam changed pull request status to closed

Sign up or log in to comment