RP3Net 0.0.1__py3-none-any.whl

RP3Net/__init__.py

@@ -0,0 +1,8 @@
+from .model import RP3Net, RP3_DEFAULT_CONFIG, RP3_CONFIG_B, load_model
+from .rp3_main import rp3_main
+import importlib
+if importlib.util.find_spec('lightning') is not None:
+    from .rp3_train import rp3_train
+else:
+    def rp3_train():
+        raise ImportError("Please install 'RP3Net[training]' to enable training")

RP3Net/fm_cfg/esm2_650m/config.json

@@ -0,0 +1,29 @@
+{
+  "architectures": [
+    "EsmForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "classifier_dropout": null,
+  "emb_layer_norm_before": false,
+  "esmfold_config": null,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 1280,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "is_folding_model": false,
+  "layer_norm_eps": 1e-05,
+  "mask_token_id": 32,
+  "max_position_embeddings": 1026,
+  "model_type": "esm",
+  "num_attention_heads": 20,
+  "num_hidden_layers": 33,
+  "pad_token_id": 1,
+  "position_embedding_type": "rotary",
+  "token_dropout": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.25.0.dev0",
+  "use_cache": true,
+  "vocab_list": null,
+  "vocab_size": 33
+}

RP3Net/fm_cfg/esm2_650m/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "<cls>",
+  "eos_token": "<eos>",
+  "mask_token": "<mask>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

RP3Net/fm_cfg/esm2_650m/tokenizer_config.json

@@ -0,0 +1,4 @@
+{
+  "model_max_length": 1000000000000000019884624838656,
+  "tokenizer_class": "EsmTokenizer"
+}

RP3Net/fm_cfg/esm2_650m/vocab.txt

@@ -0,0 +1,33 @@
+<cls>
+<pad>
+<eos>
+<unk>
+L
+A
+G
+V
+S
+E
+R
+T
+I
+D
+P
+K
+Q
+N
+F
+Y
+M
+H
+W
+C
+X
+B
+U
+Z
+O
+.
+-
+<null_1>
+<mask>

RP3Net/model/__init__.py

@@ -0,0 +1 @@
+from .model import RP3Net, load_model, Mode_Training, RP3_DEFAULT_CONFIG, RP3_CONFIG_B

RP3Net/model/layers.py

@@ -0,0 +1,171 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.nn.attention as attn
+from torch.nn import init
+import math
+# from lightning.pytorch.cli import instantiate_class
+from typing import Optional, Mapping, Sequence, Union
+# from .util import get_rank, safe_gather
+
+TNonLin = Optional[Union[str,Mapping]]
+
+def apply_p_drop(m:nn.Module, p:float):
+    def set_p_drop(m:nn.Module):
+        if isinstance(m, nn.Dropout):
+            m.p = p
+    m.apply(set_p_drop)
+
+class StackedLinear(torch.nn.Module):
+    def __init__(self, d_in:int, d_out:int, n_stack:int, bias:bool=True,
+                 device=None, dtype=None, transpose=False) -> None:
+        super().__init__()
+        self.d_in = d_in
+        self.d_out = d_out
+        self.n_stack = n_stack
+        self.transpose = transpose
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        self.weight = nn.Parameter(torch.empty((n_stack, d_in, d_out), **factory_kwargs))
+        if bias:
+            self.bias = nn.Parameter(torch.empty(n_stack, 1, d_out, **factory_kwargs))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        # Lifted from torch.nn.Linear
+        # NB: this is important; you could get really weird results (loss off the scale) if you get this wrong
+        for i in range(self.n_stack):
+            init.kaiming_uniform_(self.weight[i], a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            for i in range(self.n_stack):
+                init.uniform_(self.bias[i], -bound, bound)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = x @ self.weight        
+        if self.bias is not None:
+            out += self.bias
+        if self.transpose:
+            out = out.transpose(0,1)
+        return out
+        
+    def extra_repr(self) -> str:
+        return 'd_in={}, d_out={}, n_stack={}, bias={}'.format(
+            self.d_in, self.d_out, self.n_stack, self.bias is not None
+        )    
+
+
+class ClsHeadBlock(nn.Module):
+    def __init__(self, in_dim:int, out_dim:int, n_stack:Optional[int], layer_norm: bool, p_drop:Optional[float],
+                 nonlinearity:TNonLin, bias:bool) -> None:
+        super().__init__()
+        layer = StackedLinear(in_dim, out_dim, n_stack, bias=bias) if n_stack is not None else nn.Linear(in_dim, out_dim, bias=bias)
+        self.layers = nn.Sequential(layer)
+        if layer_norm:
+            self.layers.append(nn.LayerNorm(out_dim))
+        if p_drop is not None and p_drop != 0:
+            self.layers.append(nn.Dropout(p_drop))
+        if nonlinearity is not None:
+            nonlin = getattr(nn, nonlinearity)()
+            self.layers.append(nonlin)
+    
+    def forward(self, x):
+        return self.layers(x)
+
+
+class StackedClsHead(nn.Module):
+
+    def __init__(self, embedding_dim:int, n_stack:int, layers:Optional[Mapping[str, int]]=None,
+                 layer_norm:bool=False, p_drop:Optional[float]=None, 
+                 nonlinearity:TNonLin=None, bias:bool=False, end_bias:bool=False) -> None:
+        super().__init__()
+        layer_dims = []
+        if layers:
+            layer_dims = [round(layers['d'])] * round(layers['n'])
+        layer_dims = [embedding_dim] + layer_dims
+        self.layers = nn.Sequential(*[ClsHeadBlock(layer_dims[i], layer_dims[i+1], n_stack,
+                                                   layer_norm, p_drop, nonlinearity, bias) 
+                                    for i in range(len(layer_dims) - 1)])
+        self.layers.append(StackedLinear(layer_dims[-1], 2, n_stack, bias=end_bias, transpose=True)) # todo p_drop
+
+    def forward(self, x):
+        return self.layers(x)
+    
+
+class ClsHead(nn.Module):
+
+    def __init__(self, embedding_dim:int, layers:Optional[Mapping[str, int]]=None,
+                 layer_norm:bool=False, p_drop:Optional[float]=None, 
+                 nonlinearity:TNonLin=None, bias:bool=False, end_bias:bool=False,
+                 n_stack:Optional[int]=None, out_dim:int=2) -> None:
+        super().__init__()
+        layer_dims = []
+        if layers:
+            layer_dims = [round(layers['d'])] * round(layers['n'])
+        layer_dims = [embedding_dim] + layer_dims
+        self.layers = nn.Sequential(*[ClsHeadBlock(layer_dims[i], layer_dims[i+1], n_stack,
+                                                   layer_norm, p_drop, nonlinearity, bias) 
+                                    for i in range(len(layer_dims) - 1)])
+        self.layers.append(ClsHeadBlock(layer_dims[-1], out_dim, n_stack=None, 
+                                        layer_norm=layer_norm, p_drop=p_drop, nonlinearity=None, bias=end_bias))
+        
+    def forward(self, x):
+        return self.layers(x)
+    
+class SetTransformerPooling(nn.Module):
+    def __init__(self, d:int, num_heads:int, seq_dim:int|None=None, num_seeds:int=1, layer_norm:bool=False, p_drop:float=.0, keep_dim:bool=False,
+                 can_use_efficient:bool=True) -> None:
+        """
+        [Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks. Lee et al, ICLM 2019](https://arxiv.org/pdf/1810.00825v3.pdf)
+        Section 3.2
+        """
+        super().__init__()
+        self.seed = nn.Parameter(torch.empty(1, num_seeds, d))
+        nn.init.xavier_uniform_(self.seed)
+        if seq_dim is None:
+            seq_dim = d
+        self.mha = nn.MultiheadAttention(d, num_heads, dropout=p_drop, batch_first=True, kdim=seq_dim, vdim=seq_dim)
+        self.layer_norm = nn.LayerNorm(d) if layer_norm else None
+        self.keep_dim = keep_dim
+        self.attn_flags = [attn.SDPBackend.FLASH_ATTENTION, attn.SDPBackend.MATH, attn.SDPBackend.CUDNN_ATTENTION]
+        if can_use_efficient:
+            self.attn_flags.append(attn.SDPBackend.EFFICIENT_ATTENTION)
+
+    def forward(self, x:torch.Tensor, mask:Optional[torch.Tensor]=None, **kwargs) -> torch.Tensor:
+        if mask is not None:
+            mask = ~(mask.to(dtype=bool))
+        seed = self.seed.expand(x.shape[0], -1, -1)
+        # Required for EMLC to work.
+        # Forward mode differentiation is not implemented for memory efficient attention, 
+        # so need to disable optimization here. NB: this might change in future version of PyTorch
+        with attn.sdpa_kernel(self.attn_flags):
+            out, _ = self.mha(seed, x, x, key_padding_mask=mask, need_weights=False)
+        if self.layer_norm is not None:
+            out = self.layer_norm(out)
+        if not self.keep_dim:
+            out = out.flatten(start_dim=1)
+        return out
+    
+class MeanPooling(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x:torch.Tensor, *, mask:torch.Tensor, **kwargs) -> torch.Tensor:
+        mask = mask.to(dtype=torch.bool)
+        mask = mask.unsqueeze(-1)
+        batch_seq_len = mask.sum(dim=1)
+        x_sum = x.masked_fill(~mask, 0).sum(dim=1)
+        return x_sum / batch_seq_len
+    
+class MaxPooling(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x:torch.Tensor, *, mask:torch.Tensor, **kwargs) -> torch.Tensor:
+        mask = mask.to(dtype=torch.bool)
+        mask = mask.unsqueeze(-1)
+        x = x.masked_fill(~mask, float('-inf'))
+        ret, _ = x.max(dim=1)
+        return ret

RP3Net/model/model.py

@@ -0,0 +1,233 @@
+import abc
+import re
+import os
+import enum
+import typing
+import torch
+import torch.nn as nn
+import ml_collections as mlc
+import transformers as hub
+import pathlib
+import peft
+
+from . import layers
+from .. import util
+
+log = util.get_logger(__name__)
+
+def resolve(filename) -> pathlib.Path:
+    if type(filename) == str:
+        filename = pathlib.Path(filename)
+    if str(filename).startswith('~'):
+        filename = filename.expanduser()
+    return filename.resolve()
+
+class Mode(enum.Flag):
+    Inference = enum.auto()
+    Training_A = enum.auto()
+    Training_B = enum.auto()
+    Training_C = enum.auto()
+    Training_D = enum.auto()
+Mode_Training = Mode.Training_A | Mode.Training_B | Mode.Training_C | Mode.Training_D
+Mode_Training_FM = Mode.Training_C | Mode.Training_D
+Mode_FM = Mode_Training_FM | Mode.Inference
+Mode_Training_Aggregation = Mode.Training_B | Mode.Training_C | Mode.Training_D
+Mode_Aggregation = Mode_Training_Aggregation | Mode.Inference
+
+class RP3Net(nn.Module):
+    def __init__(self, cfg:mlc.ConfigDict) -> None:
+        super(RP3Net, self).__init__()
+        self.fm = None
+        self.mode = Mode[cfg.get('mode', 'Inference')]
+        if self.mode in Mode_FM:
+            self._init_fm(cfg, self.mode)
+            assert self.fm is not None, "Model must be initialized"
+        if self.mode in Mode_Aggregation:
+            if cfg.aggregation == 'mean':
+                self.pooling = layers.MeanPooling()
+            elif cfg.aggregation == 'max':
+                self.pooling = layers.MaxPooling()
+            elif cfg.aggregation == 'stp':
+                self.pooling = layers.SetTransformerPooling(**cfg.stp)
+            else:
+                raise ValueError(f"Aggregation type {cfg.aggregation} not supported")
+        self.cls_head = layers.ClsHead(**cfg.classification_head)
+
+    def forward(self, batch, return_repr=False):
+        if self.mode == Mode.Training_A:
+            logits = self.cls_head(batch['embeddings'])
+        elif self.mode == Mode.Training_B:
+            global_repr = self.pooling(batch['embeddings'], mask=batch['attention_mask'])
+            logits = self.cls_head(global_repr)
+        else: # Inference, Training_CD
+            seq_repr = self.sequence_representation(batch['seq'])
+            mask = self.attention_mask(batch['seq'])
+            global_repr = self.pooling(seq_repr, mask=mask)
+            logits = self.cls_head(global_repr)
+        if return_repr:
+            return logits, global_repr
+        return logits
+    
+    # def train(self, train_mode:bool=True):
+    #     if not train_mode or self.mode == Mode.Training_CD:
+    #         super().train(train_mode)
+    #     else:
+    #         if self.mode == Mode.Inference:
+    #             raise ValueError("Model is in inference mode")
+    #         elif self.mode == Mode.Training_A:
+    #             self.fm.train(False)
+    #             self.pooling.train(False)
+    #             self.cls_head.train(True)
+    #         elif self.mode == Mode.Training_B:
+    #             self.fm.train(False)
+    #             self.pooling.train(True)
+    #             self.cls_head.train(True)
+
+    @abc.abstractmethod
+    def _init_fm(self, cfg:mlc.ConfigDict, mode:Mode):
+        pass
+    
+    @abc.abstractmethod
+    def tokenize_sequences(self, sequences:typing.Sequence[str]):
+        pass
+
+    @abc.abstractmethod
+    def sequence_representation(self, batch):
+        pass
+
+    @abc.abstractmethod
+    def attention_mask(self, batch):
+        pass
+
+    @torch.no_grad()
+    def predict(self, sequences:typing.Sequence[str]|typing.Mapping[str,str], device=None):
+        is_mapping=False
+        if isinstance(sequences, typing.Mapping):
+            is_mapping = True
+            keys = list(sequences.keys())
+            sequences = [sequences[k] for k in keys]
+        seq_batch = self.tokenize_sequences(sequences)
+        if device:
+            seq_batch = seq_batch.to(device)
+        batch = {'seq': seq_batch}
+        logits = self(batch)
+        logits_norm = torch.softmax(logits, dim=-1)
+        logits_norm = logits_norm[:,1].cpu()
+        if is_mapping:
+            return {k:logits_norm[i].item() for i,k in enumerate(keys)}
+        else:
+            return logits_norm
+
+class RP3Esm2(RP3Net):
+    def _init_esm2(self, cfg:mlc.ConfigDict, mode:Mode, cfg_path:str|os.PathLike) -> None:
+        esm_cfg = hub.EsmConfig.from_pretrained(cfg_path, local_files_only=True)
+        if (checkpoint_file := cfg.get('fm.cp')):
+            log.info(f"Loading pre-trained FM from checkpoint {checkpoint_file}")
+            state_dict = torch.load(util.resolve(checkpoint_file), map_location='cpu', weights_only=True)
+            self.fm = hub.EsmModel.from_pretrained(
+                None, 
+                config=esm_cfg,
+                state_dict=state_dict, 
+                local_files_only=True, 
+                add_pooling_layer=False
+            )
+        else:
+            log.info(f"Loading random model")
+            self.fm = hub.EsmModel(esm_cfg, add_pooling_layer=False)
+        self.tokenizer = hub.EsmTokenizer.from_pretrained(cfg_path, do_lower_case=False)
+        self.re_aa_x = re.compile(r"[UZOB]")
+        lora_config=cfg.get('fm.lora')
+        if lora_config:
+            lora_config = peft.LoraConfig(**lora_config, inference_mode=(mode == Mode.Inference))
+            self.fm = peft.get_peft_model(self.fm, lora_config)
+
+    def tokenize_sequences(self, seqs:typing.Sequence[str]):
+        seqs = [self.re_aa_x.sub('X', s) for s in seqs]
+        return self.tokenizer(seqs, padding=True, return_tensors='pt')
+    
+    def sequence_representation(self, batch):
+        return self.fm(**batch).last_hidden_state
+    
+    def attention_mask(self, batch):
+        return batch['attention_mask']
+
+    def train(self, train_mode:bool=True):
+        super().train(train_mode)
+        if train_mode and self.mode in Mode_Training_FM:
+            self.fm.embeddings.position_embeddings.requires_grad_(False)
+            self.fm.contact_head.requires_grad_(False)
+        return self
+    
+    
+class RP3Esm2_650m(RP3Esm2):
+    def _init_fm(self, cfg:mlc.ConfigDict, mode:Mode):
+        cfg_path = pathlib.Path(__file__).resolve().parent.parent /'fm_cfg'/'esm2_650m'
+        self._init_esm2(cfg, mode, cfg_path)
+
+
+RP3_DEFAULT_CONFIG = mlc.FrozenConfigDict({
+    'fm': {
+        'type': 'esm2_650m',
+        'lora': {
+            'r': 8,
+            'lora_alpha': 1.0,
+            'target_modules': ['query', 'key', 'value'],
+            'lora_dropout': 0.1,
+            'bias': 'lora_only'
+        },
+    },
+    'aggregation': 'stp',
+    'stp': {
+        'seq_dim': 1280,
+        'd': 256,
+        'num_heads': 8,
+        'layer_norm': True,
+    },
+    'classification_head': {
+        'embedding_dim': 256,
+        'bias': False,
+        'end_bias': True,
+        'layer_norm': False,
+        'layers': {
+            'd': 256,
+            'n': 1
+        },
+        'nonlinearity': 'SiLU',
+    }
+})
+
+RP3_CONFIG_B = mlc.FrozenConfigDict({
+    'fm':{'type': 'esm2_650m'},
+    'aggregation': 'stp',
+    'stp': {
+        'seq_dim': 1280,
+        'd': 128,
+        'num_heads': 8,
+        'layer_norm': True,
+    },
+    'classification_head': {
+        'embedding_dim': 128,
+        'bias': False,
+        'end_bias': True,
+        'layer_norm': False,
+        'layers': {
+            'd': 512,
+            'n': 1
+        },
+        'nonlinearity': 'SiLU',
+    }
+})
+
+
+def load_model(cfg:mlc.ConfigDict, cp_path:str|os.PathLike|None=None) -> RP3Net:
+    model_type = cfg.fm.type
+    mode = Mode[cfg.get('mode', 'Inference')]
+    if model_type == 'esm2_650m':
+        model = RP3Esm2_650m(cfg)
+    else:
+        raise ValueError(f"Model {model_type} not supported")
+    if cp_path:
+        cp = torch.load(cp_path, map_location='cpu', weights_only=True)
+        model.load_state_dict(cp, strict=True)
+    return model
+        

RP3Net/rp3_main.py

@@ -0,0 +1,85 @@
+
+from typing import Optional
+import yaml
+import argparse
+import logging
+from tqdm import tqdm
+import ml_collections as mlc
+import pandas as pd
+import numpy as np
+
+import RP3Net.util as util
+import RP3Net.model as model
+
+log = util.get_logger(__file__)
+
+
+def setup_args(_args: Optional[list] = None):
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, description="""
+    Predict protein expression in E.coli from sequences. Takes in a fasta file with protein sequences and ids.
+    Outputs a .csv[.gz] file with "id,score" columns. 
+    The score is the predicted probability of expression.
+    """)
+    parser.add_argument("--log_file", help="Log file. Log output to console if set to None.")
+    parser.add_argument("--log_level", default="info", help="Log level of root logger. Appender levels are appropriately hard coded.")
+    parser.add_argument("-c", "--config", help="Model configuration in yaml, or registry entry.", default="RP3_DEFAULT_CONFIG")
+    parser.add_argument("-p", "--checkpoint", required=True, help="Model checkpoint.")
+    parser.add_argument("-f", "--fasta", required=True, help="Fasta file, possibly gzipped.")
+    parser.add_argument("-d", "--device", default='cpu', help="Device to use: cpu, cuda[:n], mps, ...")
+    parser.add_argument("-b", "--batch_size", default=8, type=int, help="Batch size. Memory consumption depends on construct sequence length and batch size. If the model runs out of SLURM/GPU memory, try reducing this parameter to make it fit.")
+    parser.add_argument("-o", "--out_file", required=True, help="Output .csv[.gz] file.")
+    parser.add_argument("--progress", action=argparse.BooleanOptionalAction, default=True, help="Show progress bar.")
+    return parser.parse_args(_args)
+
+MODEL_REGISTRY = {
+    'RP3_DEFAULT_CONFIG': model.RP3_DEFAULT_CONFIG,
+    'RP3_CONFIG_B': model.RP3_CONFIG_B,
+}
+
+def rp3_main():
+    args = setup_args()
+    util.setup_logging(save_path=args.log_file, level=args.log_level, log_console=args.log_file is None)
+
+    seq_map = util.read_fasta(args.fasta)
+    seq_lens = np.array(list(map(len, seq_map.values())))
+    log.info(f"Read {len(seq_map)} sequences from {args.fasta}. Sequence lengths (mean/std/median/min/max): "
+             f"{seq_lens.mean():.2f}/{seq_lens.std():.2f}/{np.median(seq_lens):.0f}/{seq_lens.min()}/{seq_lens.max()}")
+    
+    config_path = util.resolve(args.config)
+    if config_path.exists():
+        with open(config_path, "r") as f: 
+            config = mlc.FrozenConfigDict(yaml.load(f, Loader=yaml.FullLoader))
+    elif args.config in MODEL_REGISTRY:
+        config = MODEL_REGISTRY[args.config]
+    else:
+        raise ValueError(f"Config file {args.config} not found.")
+    log.info(f"Loading model {args.config} from checkpoint {args.checkpoint}")
+    m = model.load_model(config, args.checkpoint)
+    m = m.to(device=args.device)
+    m = m.eval()
+    log.info(f"Loaded model {m}")
+    
+    seq_keys = list(seq_map.keys())
+    def batches():
+        if args.progress:
+            tqdm_desc = args.fasta.replace('.fasta', '').replace('.gz', '')[-20:]
+            r = tqdm(range(0, len(seq_keys), args.batch_size), desc=tqdm_desc)
+        else:
+            r = range(0, len(seq_keys), args.batch_size)
+        for i in r:
+            if not args.progress:
+                log.info(f"Processing batch {i // args.batch_size + 1} of {len(seq_keys) // args.batch_size + 1}.")
+            yield {k: seq_map[k] for k in seq_keys[i:i + args.batch_size]}
+
+    ret = {}
+    for b in batches():
+        ret.update(m.predict(b, device=args.device))
+    keys, scores = zip(*ret.items())
+    out_df = pd.DataFrame({'id': keys, 'score': scores})
+    log.info(f"Writing {out_df.shape[0]} rows to {args.out_file}")
+    out_df.to_csv(args.out_file, index=False)
+
+
+if __name__ == "__main__":
+    rp3_main()
+

RP3Net/rp3_train.py

@@ -0,0 +1,18 @@
+import importlib
+from RP3Net.training.cli import RP3Cli
+from RP3Net.util import util
+
+log = util.get_logger(__name__)
+
+
+def rp3_train():
+    try:
+        assert importlib.util.find_spec('lightning') is not None, "Please install 'RP3Net[training]' to enable training"
+        cli = RP3Cli()
+    except Exception as e:
+        log.error("Top level catch", exc_info=e)
+        raise e
+    
+
+if __name__ == "__main__":
+    rp3_train()

RP3Net/training/__init__.py

@@ -0,0 +1,6 @@
+from . import lm
+from . import lm_emlc
+from . import cli
+from . import data
+from . import data_emlc
+from . import metrics