BacPT-small

BacPT-small is a bacterial proteome foundation model that contextualizes ESM2 protein embeddings using the order of proteins in a genome. It uses a 10-layer RoBERTa backbone with relative key-query position embeddings and contains 51,243,360 trainable and non-trainable model parameters.

This repository contains the inference-only weights from training checkpoint epoch 1107, the exact model configuration, and the fitted input scaler.

Input

The command accepts a protein FASTA already ordered by genomic position. It does not annotate a nucleotide genome or determine gene order. The supplied pipeline:

1. omits proteins longer than 2,250 amino acids;
2. generates 480-dimensional representations with esm2_t12_35M_UR50D, averaging layer 12 over amino-acid positions;
3. applies the fitted BacPT scaler;
4. keeps at most the first 5,000 proteins and pads shorter genomes to 5,000;
5. runs BacPT with the padding mask and CUDA float16 behavior used during the original downstream inference workflow; and
6. removes padded positions from the returned output.

Installation

pip install -r requirements.txt

Usage

Download this model repository and run:

python inference.py \
  --fasta ordered_proteins.faa \
  --output bacpt_small_embeddings.npz \
  --device cuda

The original FAIR ESM v2.0.0 code and ESM2 weights are downloaded through PyTorch Hub on first use. Add --all-layers to retain every hidden layer.

The output NPZ contains:

• protein_ids: identifiers in model order;
• last_hidden_state: contextual protein embeddings with shape [proteins, 480];
• reconstruction: reconstructed ESM2 vectors produced by the self-supervised training head, with shape [proteins, 480];
• omitted_protein_ids: proteins longer than 2,250 amino acids; and
• truncated_protein_ids: proteins beyond model position 5,000.

With --all-layers, it also contains hidden_states with shape [11, proteins, 480], including the embedding output and all 10 transformer layers.

CPU inference is supported in float32. CUDA inference uses float16 autocasting to reproduce the original inference procedure more closely.

Intended use and limitations

BacPT-small produces contextual representations for proteins in complete or draft bacterial proteomes. These representations are intended as features for research tasks such as functional, genomic-context, trait, and ecological prediction.

The model expects proteins to be provided in genomic order. It does not annotate genomes, translate nucleotide sequences, handle circular-genome rotation, or choose an ordering for unordered proteins. It was trained on bacterial genomes and has not been validated as a general model for archaeal, eukaryotic, or viral proteomes. Downstream predictions require separate task models and validation.

Paper

Bacterial proteome foundation model enhances functional prediction from enzymes to ecological interactions. Citation details will be added when the manuscript record is public.