--- license: cc-by-nc-4.0 base_model: neuralbioinfo/prokbert-mini-long library_name: pytorch pipeline_tag: feature-extraction tags: - biology - genomics - bacterial-genomics - streptococcus-pneumoniae - serotyping - novelty-detection - out-of-distribution-detection --- # 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: --- ## ⚠️ License — please read before use **These weights are released under [CC-BY-NC-4.0](https://creativecommons.org/licenses/by-nc/4.0/): non-commercial use only.** ALLCAPS is a classification head that runs on top of [ProkBERT-mini-long](https://huggingface.co/neuralbioinfo/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](#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`): ```bash 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: ```bash 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. ```bash pip install torch 'transformers<5' numpy biopython safetensors ``` ```python 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: ```bash 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/`: ```bash 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](https://github.com/bacpop/pneumococcal-serotyping/blob/main/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. ```bibtex @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 = {} } ``` ```bibtex @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](https://github.com/nbrg-ppcu/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](https://github.com/AlirezaT99) and [Sam Horsfield](https://github.com/samhorsfield96) — please open an issue on [the repository](https://github.com/bacpop/pneumococcal-serotyping/issues).