modeling_allcaps.py

"""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()