ALLCAPS — pneumococcal cps locus serotyping and novel-serotype detection

ALLCAPS takes a Streptococcus pneumoniae capsular biosynthetic locus (CBL, cps) and predicts

1. whether the sequence is a cps locus at all,
2. its serotype, and
3. whether that serotype is novel — unlike anything seen in training.

Code, pipeline and issue tracker: https://github.com/bacpop/pneumococcal-serotyping

---

⚠️ License — please read before use

These weights are released under CC-BY-NC-4.0: non-commercial use only.

ALLCAPS is a classification head that runs on top of ProkBERT-mini-long, which is itself distributed under CC-BY-NC-4.0. ProkBERT is required at inference time — every query sequence is embedded by ProkBERT before ALLCAPS sees it — so the non-commercial restriction applies to the pipeline as a whole, and we pass it through rather than advertise a freedom we cannot grant.

Note the two licenses that apply to different parts of this project:

Component	License
ALLCAPS source code (the GitHub repository, and modeling_allcaps.py here)	MIT
ALLCAPS weights (this repository)	CC-BY-NC-4.0 (inherited from ProkBERT)
ProkBERT-mini-long weights (downloaded separately at runtime)	CC-BY-NC-4.0
ProkBERT source code	MIT

If you need commercial use, you must take it up with the ProkBERT authors — our code license cannot lift their restriction. Attribution to both ALLCAPS and ProkBERT is required under CC-BY-NC; see Citation.

The base_model field above records a dependency, not a fine-tune. ALLCAPS does not modify, contain or redistribute any ProkBERT weights — ProkBERT is loaded frozen, in eval() mode, and used purely as a feature extractor. This repository holds only ALLCAPS' own parameters.

---

Model description

Input DNA is split into 4 kbp chunks with 50% overlap and embedded with frozen ProkBERT-mini-long (384-d per chunk). Those chunk embeddings are then processed by ALLCAPS' own encoder:

chunk embeddings (L, 384)
  → + learned positional embeddings
  → TransformerEncoder            (1 layer, nhead=4)
  → masked mean-pool              (padding-aware)
  → projection MLP → L2-normalised 128-d locus embedding  z
  → three linear heads:
        cbl_classifier        (2 classes: capsule y/n)
        serotype_classifier   (82 serotypes)
        genogroup_classifier  (53 genogroups — NOT TRAINED, see below)

Model class TransformerTriHeadLR. A single locus embedding z serves both classification and novelty detection.

Base encoder	neuralbioinfo/prokbert-mini-long (frozen, not included here)
Chunk size / stride	4000 bp / 50% overlap
Max input length	30,000 bp
Base embedding dim	384
Locus embedding dim	128 (L2-normalised)
Trainable parameters	2,013,833
Serotype classes	82
Genogroup classes	53 (head untrained — see below)

⚠️ The genogroup head is not trained

The architecture has three heads; the released model trains two. The genogroup head was found to add nothing to serotype or novelty performance, so it was trained with loss weight 0. It still emits logits, and those logits are meaningless — the head sits at 0.8% accuracy over 53 classes, below the 1.9% you would get by guessing.

Do not read a genogroup off this model. If you want the genogroup of a call, derive it from the predicted serotype, or use the nn_serotype the novelty detector returns (which comes from the training data, not from this head).

Novel-serotype detection

After a query is classified as a cps locus, its pooled embedding z is compared by cosine distance to every training cps embedding. The locus is called novel when the distance to its nearest training neighbour (k = 1) exceeds

T = 0.014824147932735014

which is the 95th percentile of the training leave-one-out 1-NN distances — a threshold derived without ever looking at novel data. The index shipped here holds 54,102 training locus embeddings with their serotype labels, so the report also names the closest known serotype: a novel locus is placed in a neighbourhood rather than merely rejected.

An energy score (E = −T·logsumexp(logits/T)) is available as a reference baseline (energy_summary.json carries its calibrated percentiles). Across 98 leave-one-serotype-out folds, kNN beat energy on every fold-level metric, so kNN is the deployed detector.

Distances must be computed in float64. In float32, 1 − x·y cancels below machine epsilon for the many near-identical loci in this dataset, quantising ~73% of in-distribution distances toward zero and costing ~0.015 AUROC. modeling_allcaps.py already does this; if you reimplement the detector, don't lose it.

---

Files in this repository

File	Contents
model.safetensors	The 2.0M ALLCAPS parameters
config.json	Architecture + the chunking parameters inference must match
labels.json	serotype_to_idx (82) and genogroup_to_idx (53)
knn_index.npz	Novelty index: 54,102 × 128 training embeddings, their sample keys and serotypes, and the leave-one-out distance distribution
knn_config.json	The deployed operating point: k=1, cosine, p95, and the resolved threshold
energy_summary.json	Calibrated percentiles for the energy baseline
modeling_allcaps.py	Self-contained inference — architecture, kNN detector, FASTA → calls. Loads model.safetensors, falling back to transformer_model.pth
transformer_model.pth	The same weights as a torch.save checkpoint — self-describing (model_config, serotype_to_idx), and the file the ALLCAPS CLI and the pipeline load
assets/dexB_aliA_ATCC700669.fasta	The flanking genes used to cut a cps locus out of an assembly
eval/	Classification reports, confusion matrices, CV summary, and the 98-fold novelty report

The knn_index.npz carries training contig identifiers (cbl|Public_ID#Contig_ID, the pipeline's own record keys) and their serotypes. These are public GPS/ENA accessions.

How to use

The ALLCAPS CLI (recommended)

The CLI takes whole assemblies, cuts the cps locus out for you, and downloads everything in this repository on first run (cached under ~/.cache/allcaps, override with ALLCAPS_CACHE):

git clone https://github.com/bacpop/pneumococcal-serotyping && cd ALLCAPS
mamba env create -f environment-cpu.yml      # or environment-gpu.yml on linux-64 + CUDA
conda activate allcaps && pip install .

ALLCAPS predict --input samples.txt --extract align --output results/

samples.txt lists one assembly FASTA per line (gzip is fine); bare paths work too — ALLCAPS predict a.fasta b.fasta --extract align. The run writes results/predictions.tsv and prints the same table, one row per contig, carrying serotype, serotype_confidence, is_novel_knn, nn_serotype, nn_distance and a plain-language note.

--extract align cuts between the dexB/aliA flanking genes with minimap2 — this is how the model was trained, and what you want. --extract scan embeds rolling windows across the whole assembly instead; it needs no flanking genes but is orders of magnitude slower.

To pin this repository explicitly, or to use a model you trained yourself:

ALLCAPS predict --input samples.txt --extract align --output results/ \
    --hf-repo bacpop/ALLCAPS --revision main        # or:
ALLCAPS predict --input samples.txt --extract align --output results/ \
    --model transformer_model.pth \
    --knn-index knn_index.npz \
    --energy-summary energy_summary.json \
    --offline

The three --model / --knn-index / --energy-summary files are exactly the ones in this repository, and exactly what ALLCAPS train writes into its run directory.

Standalone, without the repository

modeling_allcaps.py here is self-contained — architecture, kNN detector and FASTA → calls in one file, needing only torch, transformers, numpy, safetensors and biopython. It takes already-cut loci, not whole genomes.

pip install torch 'transformers<5' numpy biopython safetensors

from huggingface_hub import snapshot_download
from modeling_allcaps import ALLCAPS          # from this repository

path  = snapshot_download("bacpop/ALLCAPS")
model = ALLCAPS.from_pretrained(path, device="cpu")

for call in model.predict_fasta("my_loci.fasta"):
    print(call.record_id, call.serotype, call.confidence,
          "NOVEL" if call.is_novel else call.nn_serotype)

Or from the command line:

python modeling_allcaps.py my_loci.fasta --model_dir . --output calls.csv

ProkBERT is fetched from the Hub automatically on first use.

Input must be cut loci, not whole genomes. Extract the region between the flanking dexB/aliA genes first — assets/dexB_aliA_ATCC700669.fasta is the query file for that, and allcaps.data_locus_cutter in the GitHub repository does the cutting. The ALLCAPS CLI above does it for you.

The pipeline modules directly

Useful when you want the intermediate files. Run as modules from src/:

cd src

python -m allcaps.trihead.process_trihead_query \
    --query my_loci.fasta --model_path transformer_model.pth \
    --energy_summary energy_summary.json --output_dir out/ --device cpu

python -m allcaps.knn_ood predict \
    --input_type query --embeddings out/query_embeddings.npz \
    --knn_index knn_index.npz --threshold_percentile 95.0 --max_k 5 \
    --output out/knn_query_distances.csv

Output	Contents
query_results.csv	Serotype calls, confidence, is_novel_energy (the baseline)
knn_query_distances.csv	The deployed novelty call: is_novel_knn, distance, nearest known serotype

ALLCAPS predict runs exactly these two stages and joins their output into one table.

modeling_allcaps.py and the pipeline agree to within 1e-15 on distances and exactly on novelty flags. They can disagree on which neighbour is reported when several training loci sit at identical distance (~0.4% of queries, all exact ties) — the tie-break order differs, the distance does not.

---

Intended use and limitations

Intended for research and public-health genomic surveillance of S. pneumoniae: in-silico serotyping of assembled genomes, and triage of candidate novel capsular types for follow-up.

Not intended for clinical diagnosis or treatment decisions. This model is research software; it has not been validated as a diagnostic device, and the non-commercial license precludes that use in any case.

Limitations

• Trained on S. pneumoniae only — behaviour on other streptococci or other species is undefined, though the capsule y/n head is designed to reject non-cps input.
• Serotypes with fewer than 5 distinct samples were not kept as resolvable classes, so rare types may surface as "novel" rather than as themselves. 21 of the 98 serotype labels in the novelty index are represented by fewer than 20 contigs.
• Class support is very uneven: the largest serotype contributes 5,429 training contigs, the smallest 1.
• A cps locus is frequently split across two contigs in a draft assembly. Metrics here are per contig, not per sample.
• Novelty is a distance threshold, not proof. A "novel" call is a prompt to look, not a discovery.
• Training data reflects the geographic and temporal sampling of the GPS/Monocle collection.
• The genogroup head is untrained (see above).

Training data

S. pneumoniae assemblies from the GPS (Global Pneumococcal Sequencing) / Monocle collection.

Training contigs	119,909, of which 54,102 capsulated
Training assemblies with a cps locus	34,125
Held-out test contigs	13,337 (6,014 capsulated) across 7,598 assemblies
Serotype labels present	98 (82 resolvable as classes)
Collection version / date

Train/test splitting is grouped by sample, never by contig, so sibling contigs of one assembly never straddle the boundary; a runtime assertion enforces it. Serotypes with < 5 distinct samples go entirely to train, so every test serotype is also represented in training.

Trained for 100 epochs, batch size 128, AdamW at lr 1e-3, with 5-fold stratified CV followed by a refit on all data. Full hyperparameters: see TRAINING.md in the GitHub repository, or train your own with ALLCAPS train --input manifest.csv --output run/.

Evaluation

Capsule detection (training set, per contig, n = 106,909): precision, recall and F1 all 1.00; weighted F1 0.9997.

Closed-set serotyping

Split	Accuracy	Macro F1	Weighted F1	n
5-fold CV (mean ± sd)	0.952 ± 0.003	—	—	23,982/fold
Held-out GPS test	0.96	0.81	0.95	5,833
UK benchmark (external)	0.89	0.69	0.88	3,012

The macro/weighted gap is the long tail: rare serotypes carry most of the errors.

Novel-serotype detection, 98 leave-one-serotype-out folds — each fold retrains with one serotype removed, so that serotype is genuinely unseen:

Detector	AUROC mean	AUROC median	TPR@5%FPR mean	TPR@5%FPR median	folds won
kNN (deployed)	0.945	0.967	0.635	0.777	95/98
energy (baseline)	0.836	0.884	0.394	0.310	3/98

95% CI on mean AUROC: kNN [0.936, 0.955], energy [0.808, 0.864] — non-overlapping.

Threshold calibration. The p95 threshold is set on training leave-one-out distances, and it transfers: 5.00% of training loci are flagged (by construction) and 4.89% of held-out capsulated test loci — all of which are in-distribution, so that is the false-positive rate. Of held-out non-capsulated contigs, 99.9% are flagged, which is the correct behaviour: non-cps sequence is far from every cps embedding.

Per-fold and per-genogroup breakdowns are in eval/novel_detection_loo_report.md.

Citation

If you use ALLCAPS, please cite both ALLCAPS and ProkBERT.

@article{allcaps,
  title   = {ALLCAPS: a scalable and accurate deep-learning method for in silico pneumococcal serotyping},
  author  = {Tajmirriahi, Alireza and Mateeva, Teodora and Russell, Matthew and Lorenz, Oliver
             and Kim, Soo-Jeong and Lo, Stephanie and Bentley, Stephen and Lees, John
             and Horsfield, Sam},
  year    = {2026},
  journal = {bioRxiv},
  doi     = {<!-- TODO: DOI on preprint -->}
}

@article{ligeti2024prokbert,
  title   = {ProkBERT family: genomic language models for microbiome applications},
  author  = {Ligeti, Bal{\'a}zs and Szepesi-Nagy, Istv{\'a}n and Bodn{\'a}r, Bal{\'a}zs
             and Ligeti-Nagy, No{\'e}mi and Juh{\'a}sz, J{\'a}nos},
  journal = {Frontiers in Microbiology},
  volume  = {14},
  pages   = {1331233},
  year    = {2024},
  doi     = {10.3389/fmicb.2023.1331233}
}

Acknowledgements

Built on ProkBERT by the Neural Bioinformatics Research Group, Pázmány Péter Catholic University. Training data from the Global Pneumococcal Sequencing project.

Contact

Alireza Tajmirriahi and Sam Horsfield — please open an issue on the repository.