README.md

---
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: <https://github.com/bacpop/pneumococcal-serotyping>

---

## ⚠️ 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 | <!-- TODO: Monocle metadata snapshot date and accession list --> |

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     = {<!-- TODO: DOI on preprint -->}
}
```

```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).