"""Self-contained ALLCAPS inference. Everything needed to go from a *cps* locus FASTA to serotype calls and novelty flags, with no dependency on the ALLCAPS repository — only torch, transformers, numpy and (for `predict_fasta`) biopython. from modeling_allcaps import ALLCAPS model = ALLCAPS.from_pretrained(".") # or a local snapshot directory for call in model.predict_fasta("my_loci.fasta"): print(call["record_id"], call["serotype"], call["is_novel"]) This mirrors `allcaps/inference.py` + `allcaps/trihead/process_trihead_query.py` + `allcaps/knn_ood.py` in https://github.com/bacpop/pneumococcal-serotyping. The repository version is the reference implementation; this file exists so the checkpoint is usable on its own. If the two ever disagree, the repository is right. For whole assemblies rather than pre-cut loci, use the repository's `ALLCAPS predict --input samples.txt --extract align` instead — it cuts the locus out for you and fetches these artifacts from the Hub. Licence: the code is MIT, the weights are CC-BY-NC-4.0. See README.md. """ from __future__ import annotations import json import os from dataclasses import dataclass from typing import Iterator, Optional import numpy as np import torch import torch.nn as nn import torch.nn.functional as F __all__ = ["ALLCAPS", "TransformerTriHeadLR", "SerotypeCall"] # ══════════════════════════════════════════════════════════════════════════════ # Architecture — must stay byte-compatible with the published state dict # ══════════════════════════════════════════════════════════════════════════════ def _make_padding_mask(x: torch.Tensor) -> torch.Tensor: """(B, L) mask, True at all-zero (padded) chunk positions.""" return x.abs().sum(dim=-1) == 0 def _masked_mean_pool(x: torch.Tensor, padding_mask: torch.Tensor) -> torch.Tensor: """Mean over real positions only. Averaging over padding too would make a locus's embedding depend on what else happened to be in its batch — the pooled vector feeds the novelty detector, so that would move the distances. """ real = (~padding_mask).float().unsqueeze(-1) lengths = real.sum(dim=1).clamp(min=1) return (x * real).sum(dim=1) / lengths class TransformerTriHeadLR(nn.Module): """Chunk embeddings → transformer encoder → pooled 128-d locus embedding → 3 heads. Head 3 (`genogroup_classifier`) is present but was trained with weight 0 in the released checkpoint; its logits are meaningless. See README.md. """ def __init__(self, input_dim, num_classes, output_dim=128, max_len=64, nhead=4, num_layers=2, **kwargs): super().__init__() self.pos_embed = nn.Embedding(max_len, input_dim) encoder_layer = nn.TransformerEncoderLayer( d_model=input_dim, nhead=nhead, dim_feedforward=4 * input_dim, batch_first=True, ) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) self.cbl_classifier = nn.Linear(output_dim, 2) self.serotype_classifier = nn.Linear(output_dim, num_classes[0]) self.genogroup_classifier = nn.Linear(output_dim, num_classes[1]) self.project = nn.Sequential( nn.Linear(input_dim, input_dim), nn.ReLU(), nn.Linear(input_dim, output_dim) ) def forward(self, x): B, L, D = x.size() padding_mask = _make_padding_mask(x) pos = torch.arange(L, device=x.device).unsqueeze(0) x = x + self.pos_embed(pos) x = self.encoder(x, src_key_padding_mask=padding_mask) x = _masked_mean_pool(x, padding_mask) z = F.normalize(self.project(x), dim=1) return self.cbl_classifier(z), (self.serotype_classifier(z), self.genogroup_classifier(z)), z # ══════════════════════════════════════════════════════════════════════════════ # Novelty detector # ══════════════════════════════════════════════════════════════════════════════ class KnnNoveltyIndex: """Cosine distance to the k-th nearest training locus, thresholded. Brute force over ~54k x 128 — a few hundred ms per query batch on CPU, so there is no reason to carry an approximate index. """ def __init__(self, embeddings, serotypes, keys, threshold, k=1, metric="cosine"): if metric != "cosine": raise ValueError(f"only cosine is supported, got {metric!r}") # float64 throughout. In float32 the `1 - x.y` of cosine distance 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. This is not a micro-optimisation to undo. X = np.ascontiguousarray(embeddings, dtype=np.float64) self.X = X / np.linalg.norm(X, axis=1, keepdims=True) self.serotypes = np.asarray(serotypes) self.keys = np.asarray(keys) self.threshold = float(threshold) self.k = int(k) @classmethod def from_pretrained(cls, path: str): with np.load(os.path.join(path, "knn_index.npz"), allow_pickle=False) as npz: embeddings = npz["train_data"] serotypes = npz["train_serotypes"] keys = npz["train_keys"] with open(os.path.join(path, "knn_config.json")) as f: cfg = json.load(f) return cls(embeddings, serotypes, keys, threshold=cfg["threshold"], k=cfg["k"], metric=cfg["metric"]) def query(self, z: np.ndarray, top_k: int = 1): """Returns (distances, serotypes, keys), each (n_queries, top_k), nearest first.""" Q = np.ascontiguousarray(z, dtype=np.float64) Q = Q / np.linalg.norm(Q, axis=1, keepdims=True) distances = 1.0 - Q @ self.X.T top_k = min(top_k, distances.shape[1]) idx = np.argpartition(distances, top_k - 1, axis=1)[:, :top_k] order = np.take_along_axis(distances, idx, axis=1).argsort(axis=1) idx = np.take_along_axis(idx, order, axis=1) return (np.take_along_axis(distances, idx, axis=1), self.serotypes[idx], self.keys[idx]) def is_novel(self, distance_to_kth: np.ndarray) -> np.ndarray: return distance_to_kth > self.threshold # ══════════════════════════════════════════════════════════════════════════════ # End-to-end model # ══════════════════════════════════════════════════════════════════════════════ @dataclass class SerotypeCall: record_id: str is_cps: bool cps_probability: float serotype: Optional[str] confidence: float is_novel: Optional[bool] nn_distance: Optional[float] nn_serotype: Optional[str] embedding: np.ndarray def as_dict(self, with_embedding: bool = False) -> dict: d = {k: v for k, v in self.__dict__.items() if k != "embedding"} if with_embedding: d["embedding"] = self.embedding return d # Mapping-style access so callers can treat these as plain records. def __getitem__(self, key): return getattr(self, key) class ALLCAPS: """ProkBERT → TriHead encoder → serotype call + novelty flag.""" def __init__(self, head, config, serotype_to_idx, genogroup_to_idx, knn: Optional[KnnNoveltyIndex] = None, device="cpu", base_model_name="neuralbioinfo/prokbert-mini-long"): self.head = head self.config = config self.device = torch.device(device) self.idx_to_serotype = {v: k for k, v in serotype_to_idx.items()} self.idx_to_genogroup = {v: k for k, v in genogroup_to_idx.items()} self.knn = knn self.base_model_name = base_model_name self._tokenizer = None self._base = None @classmethod def from_pretrained(cls, path: str = ".", device: str = "cpu", load_knn: bool = True) -> "ALLCAPS": """Load from a directory holding config.json, labels.json and the weights. Pass a `huggingface_hub.snapshot_download(...)` path to load from the Hub. """ with open(os.path.join(path, "config.json")) as f: config = json.load(f) with open(os.path.join(path, "labels.json")) as f: labels = json.load(f) head = TransformerTriHeadLR( input_dim=config["input_dim"], num_classes=tuple(config["num_classes"]), output_dim=config["output_dim"], max_len=config.get("max_len", 64), nhead=config["nhead"], num_layers=config["num_layers"], ) safetensors_path = os.path.join(path, "model.safetensors") if os.path.exists(safetensors_path): from safetensors.torch import load_file state = load_file(safetensors_path) else: # the original torch checkpoint carries the same weights state = torch.load(os.path.join(path, "transformer_model.pth"), map_location="cpu")["model_state_dict"] head.load_state_dict(state) head.eval().to(device) knn = None if load_knn and os.path.exists(os.path.join(path, "knn_index.npz")): knn = KnnNoveltyIndex.from_pretrained(path) return cls(head, config, labels["serotype_to_idx"], labels["genogroup_to_idx"], knn=knn, device=device, base_model_name=config.get("base_model", "neuralbioinfo/prokbert-mini-long")) # ── ProkBERT, loaded lazily so `from_pretrained` stays cheap ────────────── def _ensure_base_model(self): if self._base is not None: return from transformers import AutoModel, AutoTokenizer self._tokenizer = AutoTokenizer.from_pretrained( self.base_model_name, trust_remote_code=True) self._base = AutoModel.from_pretrained( self.base_model_name, trust_remote_code=True).to(self.device) self._base.eval() @torch.no_grad() def embed_chunks(self, sequence: str) -> torch.Tensor: """DNA string → (1, n_chunks, 384) ProkBERT chunk embeddings.""" self._ensure_base_model() chunk_size = min(self.config["chunk_size"], getattr(self._tokenizer, "model_max_length", self.config["chunk_size"])) stride = int(chunk_size * self.config["stride_ratio"]) sequence = sequence[: self.config["max_length"]].upper() chunks = [sequence[i:i + chunk_size] for i in range(0, max(len(sequence) - chunk_size + 1, 1), stride)] chunks = [c for c in chunks if len(c) > 0] or [sequence] vectors = [] for chunk in chunks: enc = self._tokenizer(chunk, return_tensors="pt", truncation=True, max_length=chunk_size).to(self.device) out = self._base(**enc).last_hidden_state # (1, tokens, 384) vectors.append(out.mean(dim=1).squeeze(0)) return torch.stack(vectors).unsqueeze(0) # (1, n_chunks, 384) @torch.no_grad() def predict_sequence(self, sequence: str, record_id: str = "query", top_k: int = 1) -> SerotypeCall: chunk_embeddings = self.embed_chunks(sequence).to(self.device) cps_logits, (serotype_logits, _), z = self.head(chunk_embeddings) cps_probs = torch.softmax(cps_logits, dim=-1).squeeze(0).cpu().numpy() is_cps = bool(cps_probs[1] > self.config.get("capsule_threshold", 0.5)) serotype_probs = torch.softmax(serotype_logits, dim=-1).squeeze(0).cpu().numpy() best = int(serotype_probs.argmax()) z_np = z.squeeze(0).cpu().numpy() is_novel = nn_distance = nn_serotype = None if self.knn is not None: distances, serotypes, _ = self.knn.query(z_np[None, :], top_k=max(top_k, self.knn.k)) nn_distance = float(distances[0, 0]) nn_serotype = str(serotypes[0, 0]) # The score is the distance to the k-th neighbour; k=1 as deployed. kth = float(distances[0, min(self.knn.k, distances.shape[1]) - 1]) is_novel = bool(kth > self.knn.threshold) return SerotypeCall( record_id=record_id, is_cps=is_cps, cps_probability=float(cps_probs[1]), # A non-cps sequence has no serotype; reporting the argmax anyway is # how people end up quoting a serotype for a housekeeping contig. serotype=self.idx_to_serotype[best] if is_cps else None, confidence=float(serotype_probs[best]), is_novel=is_novel if is_cps else None, nn_distance=nn_distance, nn_serotype=nn_serotype, embedding=z_np, ) def predict_fasta(self, path: str, top_k: int = 1) -> Iterator[SerotypeCall]: """Yield one call per FASTA record. Records should be cut *cps* loci.""" from Bio import SeqIO for record in SeqIO.parse(path, "fasta"): yield self.predict_sequence(str(record.seq), record_id=record.id, top_k=top_k) if __name__ == "__main__": import argparse import csv import sys p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) p.add_argument("fasta", help="FASTA of cps loci") p.add_argument("--model_dir", default=".", help="Directory holding config.json etc.") p.add_argument("--device", default="cpu") p.add_argument("--output", default="-", help="CSV output path ('-' for stdout)") args = p.parse_args() model = ALLCAPS.from_pretrained(args.model_dir, device=args.device) handle = sys.stdout if args.output == "-" else open(args.output, "w", newline="") writer = csv.DictWriter(handle, fieldnames=[ "record_id", "is_cps", "cps_probability", "serotype", "confidence", "is_novel", "nn_distance", "nn_serotype"]) writer.writeheader() for call in model.predict_fasta(args.fasta): writer.writerow(call.as_dict()) if handle is not sys.stdout: handle.close()