flash-abb 0.0.1__py3-none-any.whl

flash_abb/__init__.py

@@ -0,0 +1,2 @@
+from .pretrained import pretrained
+from .pretrained_tap import pretrained_sss, pretrained_tap

flash_abb/load_model.py

@@ -0,0 +1,75 @@
+import os, subprocess, json, argparse,requests
+from yaml import load, Loader
+import torch
+
+list_of_models = {
+    "flash-abb":"flabb_weights.pt",
+    "flash-abb_masked":"flabb_masked_weights.pt",
+}
+flash_abb_models = ["flash-abb", "flash-abb_masked"]
+
+
+def load_model(model_to_use="flash-abb", random_init=False, device='cpu'):
+
+    if model_to_use in flash_abb_models:
+        flabb, hparams = fetch_flash_abb(
+            model_to_use, 
+            random_init=random_init, 
+            device=device
+        )
+    else: 
+        assert False, f"The selected model to use ({model_to_use}) does not exist.\
+        Please select a valid model."   
+
+    return flabb, hparams
+
+
+def fetch_flash_abb(model_to_use, random_init=False, device='cpu'):
+
+    from .model.flash_abb import FlashABB
+
+    local_model_folder = os.path.join(os.path.dirname(__file__), "weights")
+    file_model = list_of_models[model_to_use]
+
+    with open(os.path.join(local_model_folder, 'params.yaml'), 'r', encoding='utf-8') as f:
+        hparams = argparse.Namespace(**load(f, Loader=Loader)).model
+
+    flabb = FlashABB(hparams)
+    if not random_init:
+        ckpt = torch.load(
+            os.path.join(local_model_folder, file_model),
+            map_location=torch.device(device),
+            weights_only=False,
+        )
+        flabb.load_state_dict(ckpt)
+
+    return flabb, hparams
+
+
+def fetch_sss(random_init=False, device='cpu'):
+    from .model.seq2struct2seq import BERTCoords
+
+    model = BERTCoords(device=device)
+    if not random_init:
+        weights_path = os.path.join(os.path.dirname(__file__), "weights", "sss_weights.pt")
+        ckpt = torch.load(weights_path, map_location=torch.device(device), weights_only=False)
+        model.load_state_dict(ckpt)
+    return model.to(device)
+
+
+def fetch_tap(random_init=False, device='cpu'):
+    from .model.seq2struct2seq import BERTCoords
+    from .model.tap_head import TAPHead
+
+    encoder = BERTCoords(device=device)
+    head = TAPHead()
+    if not random_init:
+        weights_path = os.path.join(os.path.dirname(__file__), "weights", "tap_weights.pt")
+        ckpt = torch.load(weights_path, map_location=torch.device(device), weights_only=False)
+        # encoder_state keys have a 'model.' prefix from the training wrapper
+        encoder_state = {k.removeprefix('model.'): v for k, v in ckpt['encoder_state'].items()}
+        encoder.load_state_dict(encoder_state, strict=False)
+        head.load_state_dict(ckpt['head_state'], strict=False)
+        head.tgt_mean.copy_(ckpt['tgt_mean'])
+        head.tgt_std.copy_(ckpt['tgt_std'])
+    return encoder.to(device), head.to(device)

flash_abb/pretrained.py

@@ -0,0 +1,45 @@
+import numpy as np
+import torch
+
+from .load_model import load_model
+from .model.flash_abb import featurize, FlashABBResult
+
+
+class pretrained:
+
+    def __init__(self, model_to_use="flash-abb", random_init=False, device='cuda'):
+        super().__init__()
+        
+        self.used_device = torch.device(device)
+
+        self.flabb, self.hparams = load_model(model_to_use, random_init=random_init)
+        self.flabb.to(self.used_device)
+        self.flabb.eval() # Default
+        self.device = torch.device(device)
+
+    def freeze(self):
+        self.flabb.eval()
+
+    def unfreeze(self):
+        self.flabb.train()
+
+    def from_features(self, features, batch_size=50):
+        pred = self.flabb.model(
+            {'single': features['single']},
+            features['aatype'],
+            features['res_idx'],
+            features['mask']
+        )
+        result = FlashABBResult(seqs, pred, features['mask'])
+        return result
+
+    def __call__(self, seqs, batch_size=50):
+        features = featurize(seqs, self.device)
+        pred = self.flabb.model(
+            {'single': features['single']},
+            features['aatype'],
+            features['res_idx'],
+            features['mask']
+        )
+        result = FlashABBResult(seqs, pred, features['mask'])
+        return result

flash_abb/pretrained_tap.py

@@ -0,0 +1,167 @@
+import os
+import numpy as np
+import torch
+from .load_model import fetch_sss, fetch_tap
+from .model.tokenizer import ABtokenizer
+from .model.flag_calibrator import FlagCalibrator
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+_CALIBRATOR_PATH = os.path.join(os.path.dirname(__file__), 'weights', 'flag_calibrators.pkl')
+
+
+class SSSResult:
+    """Result from FlashABB-SSS: per-residue structure-aware embeddings."""
+
+    def __init__(self, embeddings: torch.Tensor, mask: torch.Tensor):
+        self._embeddings = embeddings
+        self._mask = mask
+
+    @property
+    def embeddings(self) -> torch.Tensor:
+        """(batch, seq_len, emb_size) per-residue embeddings."""
+        return self._embeddings
+
+    @property
+    def mask(self) -> torch.Tensor:
+        """(batch, seq_len) bool mask — True where residue is present."""
+        return self._mask
+
+
+class TAPResult:
+    """Result from FlashTAP: four antibody developability scores and flag probabilities."""
+
+    TAP_COLS = ['PSH', 'PPC', 'PNC', 'SFvCSP']
+
+    def __init__(self, tensor: torch.Tensor, flag_probs_array: np.ndarray | None = None):
+        self._tensor = tensor
+        self._flag_probs_array = flag_probs_array  # (batch, 4) or None
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        """(batch, 4) raw score tensor."""
+        return self._tensor
+
+    @property
+    def scores(self) -> list[dict]:
+        """List of dicts (one per antibody) mapping property name → float."""
+        return [
+            {col: self._tensor[i, j].item() for j, col in enumerate(self.TAP_COLS)}
+            for i in range(self._tensor.shape[0])
+        ]
+
+    @property
+    def flag_probs(self) -> list[dict] | None:
+        """List of dicts (one per antibody) mapping property name → P(flag).
+
+        Returns None if no calibrator was loaded.
+        """
+        if self._flag_probs_array is None:
+            return None
+        return [
+            {col: float(self._flag_probs_array[i, j]) for j, col in enumerate(self.TAP_COLS)}
+            for i in range(self._flag_probs_array.shape[0])
+        ]
+
+    @property
+    def any_flag_prob(self) -> list[float] | None:
+        """P(any flag) for each antibody, assuming property independence.
+
+        Returns None if no calibrator was loaded.
+        """
+        if self._flag_probs_array is None:
+            return None
+        any_flag = 1 - np.prod(1 - self._flag_probs_array, axis=1)
+        return any_flag.tolist()
+
+
+def _tokenize(seqs, alphabet: ABtokenizer, device):
+    tokens = alphabet(seqs, pad=True, w_extra_tkns=False)
+    return tokens.to(device)
+
+
+def _emb_and_mask(model, seqs, tokens, alphabet, device):
+    """Run BERTCoords forward and return (embeddings, mask) with sep removed."""
+    pad_mask = tokens.eq(alphabet.pad_token).to(device)
+    emb = model(seqs, tokens, pad_mask, return_emb=True)
+
+    sep_mask = tokens != alphabet.sep_token
+    src_shape = list(tokens.shape)
+    src_shape[1] -= 1
+    mask = (~pad_mask)[sep_mask].view(src_shape)
+    return emb, mask
+
+
+class pretrained_sss:
+    """FlashABB-SSS: structure-aware antibody sequence encoder.
+
+    Usage::
+
+        from flash_abb import pretrained_sss
+        sss = pretrained_sss()
+        result = sss(['EVQL...|DIQL...'])
+        print(result.embeddings.shape)   # (1, seq_len, 128)
+    """
+
+    def __init__(self, random_init: bool = False, device=DEVICE):
+        self.device = device
+        self.sss = fetch_sss(random_init=random_init, device=str(device))
+        self.sss.eval()
+        self.sss.requires_grad_(False)
+        self.alphabet = self.sss.alphabet
+
+    def __call__(self, seqs, batch_size: int = 50) -> SSSResult:
+        all_emb, all_mask = [], []
+        for i in range(0, len(seqs), batch_size):
+            batch = seqs[i:i + batch_size]
+            tokens = _tokenize(batch, self.alphabet, self.device)
+            with torch.no_grad():
+                emb, mask = _emb_and_mask(self.sss, batch, tokens, self.alphabet, self.device)
+            all_emb.append(emb)
+            all_mask.append(mask)
+        return SSSResult(torch.cat(all_emb), torch.cat(all_mask))
+
+
+class pretrained_tap:
+    """FlashTAP: predicts four TAP developability scores from antibody sequences.
+
+    Scores: PSH (patches of surface hydrophobicity), PPC (positive patches),
+    PNC (negative patches), SFvCSP (structural Fv charge symmetry parameter).
+
+    Usage::
+
+        from flash_abb import pretrained_tap
+        tap = pretrained_tap()
+        result = tap(['EVQL...|DIQL...'])
+        print(result.scores)        # [{'PSH': ..., 'PPC': ..., 'PNC': ..., 'SFvCSP': ...}]
+        print(result.flag_probs)    # [{'PSH': 0.12, 'PPC': 0.03, 'PNC': 0.05, 'SFvCSP': 0.41}]
+        print(result.any_flag_prob) # [0.52]
+    """
+
+    def __init__(self, random_init: bool = False, device=DEVICE):
+        self.device = device
+        self.encoder, self.head = fetch_tap(random_init=random_init, device=str(device))
+        self.encoder.eval()
+        self.encoder.requires_grad_(False)
+        self.head.eval()
+        self.head.requires_grad_(False)
+        self.alphabet = self.encoder.alphabet
+        self.calibrator = FlagCalibrator.load(_CALIBRATOR_PATH) if not random_init else None
+
+    def __call__(self, seqs, batch_size: int = 50) -> TAPResult:
+        all_scores = []
+        for i in range(0, len(seqs), batch_size):
+            batch = seqs[i:i + batch_size]
+            tokens = _tokenize(batch, self.alphabet, self.device)
+            with torch.no_grad():
+                emb, mask = _emb_and_mask(self.encoder, batch, tokens, self.alphabet, self.device)
+                scores = self.head(emb, mask)
+            all_scores.append(scores)
+        score_tensor = torch.cat(all_scores)
+
+        flag_probs_array = None
+        if self.calibrator is not None:
+            scores_np = score_tensor.cpu().numpy()
+            flag_probs_array = self.calibrator.predict_proba(scores_np)
+
+        return TAPResult(score_tensor, flag_probs_array)

flash_abb-0.0.1.dist-info/METADATA

@@ -0,0 +1,105 @@
+Metadata-Version: 2.4
+Name: flash-abb
+Version: 0.0.1
+Summary: Flash-ABB: modelling antibody structures at the speed of language
+Home-page: https://github.com/oxpig/FlashABB
+Author: Isaac Ellmen
+Maintainer: Isaac Ellmen
+Maintainer-email: isaac.ellmen@stats.ox.ac.uk
+License: BSD 3-clause license
+Description-Content-Type: text/markdown
+Requires-Dist: torch>2
+Requires-Dist: requests
+Requires-Dist: einops
+Requires-Dist: rotary-embedding-torch
+Requires-Dist: ml_collections
+Requires-Dist: numpy
+Requires-Dist: dm-tree
+Requires-Dist: pyyaml
+Requires-Dist: scipy
+Dynamic: author
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: maintainer
+Dynamic: maintainer-email
+Dynamic: requires-dist
+Dynamic: summary
+
+# FlashABB: modelling antibody structures at the speed of language
+
+![Inference speed comparison](figures/speedup_multiplier.png)
+
+Installation:
+
+PyPi coming soon
+
+```bash
+git clone git@github.com:oxpig/FlashABB.git
+cd FlashABB
+pip install .
+```
+
+## Structure prediction
+
+The following is also in `example.py` and can be used to create the structures in `sample_preds`.
+
+```python
+from flash_abb import pretrained
+import torch
+
+flabb = pretrained(device='cuda')
+
+# Sequences in heavy|light format
+seqs = [
+    'EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVSS|DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK',
+    'EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVS|DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK',
+]
+
+with torch.no_grad():
+    result = flabb(seqs)
+
+print(result.coords.shape)          # (2, n_residues, 37, 3)
+print(result.bb_coords.shape)       # (2, n_residues, 4, 3)
+
+result.to_pdbs(['ab1', 'ab2'], pdb_dir='sample_preds')
+```
+
+## Developability scoring (FlashTAP)
+
+FlashTAP predicts four [TAP](https://doi.org/10.1038/s42003-023-05744-8) developability scores: PSH, PPC, PNC, and SFvCSP.
+
+```python
+from flash_abb import pretrained_tap
+
+tap = pretrained_tap(device='cuda')
+
+seqs = [
+    'EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVSS|DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK',
+]
+
+result = tap(seqs)
+print(result.scores)        # [{'PSH': ..., 'PPC': ..., 'PNC': ..., 'SFvCSP': ...}]
+print(result.tensor)        # (1, 4) raw score tensor
+print(result.flag_probs)    # [{'PSH': 0.12, 'PPC': 0.03, 'PNC': 0.05, 'SFvCSP': 0.41}]
+print(result.any_flag_prob) # [0.47]
+```
+
+## Structure-aware embeddings (FlashABB-SSS)
+
+FlashABB-SSS (seq2struct2seq) produces per-residue embeddings that combine sequence and predicted 3D structure. These can be used as features for downstream tasks.
+
+```python
+from flash_abb import pretrained_sss
+
+sss = pretrained_sss(device='cuda')
+
+seqs = [
+    'EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVSS|DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK',
+]
+
+result = sss(seqs)
+print(result.embeddings.shape)  # (1, n_residues, 128)
+print(result.mask.shape)        # (1, n_residues)
+```

flash_abb-0.0.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+flash_abb/__init__.py,sha256=6NnMNvsu4FTcWyBtklZvyStkn3HYxxueSt_7LS4faDI,94
+flash_abb/load_model.py,sha256=RaY6J1hfs3Q9N1Iy6PN-klMg0yei_NSRwrTq3gO35I0,2626
+flash_abb/pretrained.py,sha256=zPgDi2Kk1u8zpKk6ywJxL12TgWLbPJTbge--8Wg8iQw,1297
+flash_abb/pretrained_tap.py,sha256=ZhXKDBPhP2IvLM8W2q6xZ-NKYds10E-BNt2P_vaH29c,5983
+flash_abb-0.0.1.dist-info/METADATA,sha256=6RxPsHdWYW8TvUeAnN97tjlsyr5E6oG9o0Iw1Wg21TE,3441
+flash_abb-0.0.1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+flash_abb-0.0.1.dist-info/top_level.txt,sha256=cXv5m3fquiDHdQpJpR811w7tPuTRUn4WGipZG0cFxkw,10
+flash_abb-0.0.1.dist-info/RECORD,,

flash_abb-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flash_abb-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+flash_abb