pCoMole/peptide/objectives.py

import pdb
import torch
import torch.nn.functional as F
import torch.nn as nn
import pytorch_lightning as pl
from transformers import AutoModel, AutoConfig, AutoTokenizer
import xgboost as xgb

import sys
sys.path.append('/scratch/pranamlab/tong/AReUReDi/PeptiVerse/')
from inference import PeptiVersePredictor

pred = PeptiVersePredictor(
    manifest_path="/scratch/pranamlab/tong/AReUReDi/PeptiVerse/best_models.txt",          # best model list
    classifier_weight_root="/scratch/pranamlab/tong/AReUReDi/PeptiVerse/",               # repo root (where training_classifiers/ lives)
    device="cuda",                            # or "cpu"
)

from modules.bindevaluator_modules import *

class BindEvaluator(pl.LightningModule):
    def __init__(self, n_layers, d_model, d_hidden, n_head,
                 d_k, d_v, d_inner, dropout=0.2,
                 learning_rate=0.00001, max_epochs=15, kl_weight=1):
        super(BindEvaluator, self).__init__()

        self.esm_model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D")
        self.esm_model.eval()
        # freeze all the esm_model parameters
        for param in self.esm_model.parameters():
            param.requires_grad = False

        self.repeated_module = RepeatedModule3(n_layers, d_model, d_hidden,
                                               n_head, d_k, d_v, d_inner, dropout=dropout)

        self.final_attention_layer = MultiHeadAttentionSequence(n_head, d_model,
                                                                d_k, d_v, dropout=dropout)

        self.final_ffn = FFN(d_model, d_inner, dropout=dropout)

        self.output_projection_prot = nn.Linear(d_model, 1)

        self.learning_rate = learning_rate
        self.max_epochs = max_epochs
        self.kl_weight = kl_weight

        self.classification_threshold = nn.Parameter(torch.tensor(0.5))  # Initial threshold
        self.historical_memory = 0.9
        self.class_weights = torch.tensor([3.000471363174231, 0.5999811490272925])  # binding_site weights, non-bidning site weights

    def forward(self, binder_tokens, target_tokens):
        peptide_sequence = self.esm_model(**binder_tokens).last_hidden_state
        protein_sequence = self.esm_model(**target_tokens).last_hidden_state

        prot_enc, sequence_enc, sequence_attention_list, prot_attention_list, \
            seq_prot_attention_list, seq_prot_attention_list = self.repeated_module(peptide_sequence,
                                                                                    protein_sequence)

        prot_enc, final_prot_seq_attention = self.final_attention_layer(prot_enc, sequence_enc, sequence_enc)

        prot_enc = self.final_ffn(prot_enc)

        prot_enc = self.output_projection_prot(prot_enc)

        return prot_enc

    def get_probs(self, x_t, target_sequence):
        '''
        Inputs:
        - xt: Shape (bsz, seq_len)
        - target_sequence: Shape (1, tgt_len)
        '''
        # pdb.set_trace()
        target_sequence = target_sequence.repeat(x_t.shape[0], 1)
        binder_attention_mask = torch.ones_like(x_t)
        target_attention_mask = torch.ones_like(target_sequence)

        binder_attention_mask[:, 0] = binder_attention_mask[:, -1] = 0
        target_attention_mask[:, 0] = target_attention_mask[:, -1] = 0

        binder_tokens = {'input_ids': x_t, 'attention_mask': binder_attention_mask.to(x_t.device)}
        target_tokens = {'input_ids': target_sequence, 'attention_mask': target_attention_mask.to(target_sequence.device)}
        
        logits = self.forward(binder_tokens, target_tokens).squeeze(-1)
        # pdb.set_trace()
        logits[:, 0] = logits[:, -1] = -100 # float('-inf')
        probs = torch.sigmoid(logits)

        return probs    # shape (bsz, tgt_len)

    def motif_score(self, x_t, target_sequence, motifs):
        probs = self.get_probs(x_t, target_sequence)
        motif_probs = probs[:, motifs]
        motif_score = motif_probs.sum(dim=-1) / len(motifs)
        # pdb.set_trace()
        return motif_score

    def non_motif_score(self, x_t, target_sequence, motifs):
        probs = self.get_probs(x_t, target_sequence)
        non_motif_probs = probs[:, [i for i in range(probs.shape[1]) if i not in motifs]]
        mask = non_motif_probs >= 0.5
        count = mask.sum(dim=-1)

        non_motif_score = torch.where(count > 0, (non_motif_probs * mask).sum(dim=-1) / count, torch.zeros_like(count))

        return non_motif_score

    def scoring(self, x_t, target_sequence, motifs, penalty=False):
        probs = self.get_probs(x_t, target_sequence)
        motif_probs = probs[:, motifs]
        motif_score = motif_probs.sum(dim=-1) / len(motifs)
        # pdb.set_trace()

        if penalty:
            non_motif_probs = probs[:, [i for i in range(probs.shape[1]) if i not in motifs]]
            mask = non_motif_probs >= 0.5
            count = mask.sum(dim=-1)
            # non_motif_score = 1 - torch.where(count > 0, (non_motif_probs * mask).sum(dim=-1) / count, torch.zeros_like(count))
            non_motif_score = count / target_sequence.shape[1]
            return motif_score, 1 - non_motif_score
        else:
            return motif_score

class MotifModel(nn.Module):
    def __init__(self, bindevaluator, target_sequence, motifs, tokenizer, device, penalty=False):
        super(MotifModel, self).__init__()
        self.bindevaluator = bindevaluator
        self.target_sequence = target_sequence
        self.motifs = motifs
        self.penalty = penalty
        self.tokenizer = tokenizer
        self.device = device
    
    def forward(self, input_seqs):
        if self.penalty:
            specificity = []
        scores = []

        for seq in input_seqs:
            seq = self.tokenizer(seq, return_tensors='pt').to(self.device)
            x = seq['input_ids']

            score = self.bindevaluator.scoring(x, self.target_sequence, self.motifs, self.penalty)
            if self.penalty:
                specificity.append(score[1].item())
            scores.append(score[0].item())

        if self.penalty:
            return scores, specificity
        return scores

class Hemolysis:
    def __call__(self, input_seqs):
        scores = []
        for seq in input_seqs:
            score = pred.predict_property("hemolysis", col="wt", input_str=seq)['score']
            scores.append(1 - score)
        
        return torch.tensor(scores)
    
class NonFouling:
    def __call__(self, input_seqs):
        scores = []
        for seq in input_seqs:
            score = pred.predict_property("nf", col="wt", input_str=seq)['score']
            scores.append(score)
        
        return torch.tensor(scores)
    
class Solubility:
    def __init__(self):
        self.hydrophobic = list("AVLIMFWPavilmfwpŶƘṂŁĊ")
    
    def __call__(self, aa_seqs: list):
        scores = []
        for seq in aa_seqs:
            score = len([tok for tok in seq if tok not in self.hydrophobic]) / len(seq)
            scores.append(score)
        return torch.tensor(scores)
    
class Permeability:
    def __call__(self, input_seqs):
        scores = []
        for seq in input_seqs:
            score = pred.predict_property("permeability_penetrance", col="wt", input_str=seq)['score']
            # score = (score + 9) / (-4 + 9)
            scores.append(score)
        
        return torch.tensor(scores)

class HalfLife:
    def __call__(self, input_seqs):
        scores = []
        for seq in input_seqs:
            score = pred.predict_property("halflife", col="wt", input_str=seq)['score']
            scores.append(score)
        
        return torch.tensor(scores)
    
class Affinity:
    def __init__(self, target):
        self.target = target

    def __call__(self, input_seqs):
        scores = []
        for seq in input_seqs:
            score = pred.predict_binding_affinity(col="wt", target_seq=self.target, binder_str=seq)['affinity']
            scores.append(score / 10)

        return torch.tensor(scores)

def load_bindevaluator(checkpoint_path, device):
    bindevaluator = BindEvaluator.load_from_checkpoint(checkpoint_path, n_layers=8, d_model=128, d_hidden=128, n_head=8, d_k=64, d_v=128, d_inner=64).to(device)
    bindevaluator.eval()
    for param in bindevaluator.parameters():
        param.requires_grad = False

    return bindevaluator