| import pdb |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| from .layers import * |
| from .modules import * |
| import pdb |
| from transformers import EsmModel, EsmTokenizer |
|
|
| def to_var(x): |
| if torch.cuda.is_available(): |
| x = x.cuda() |
| return x |
|
|
|
|
| class RepeatedModule3(nn.Module): |
| def __init__(self, n_layers, d_model, d_hidden, |
| n_head, d_k, d_v, d_inner, dropout=0.1): |
| super().__init__() |
|
|
| self.linear1 = nn.Linear(1280, d_model) |
| self.linear2 = nn.Linear(1280, d_model) |
| self.sequence_embedding = nn.Embedding(20, d_model) |
| self.d_model = d_model |
|
|
| self.reciprocal_layer_stack = nn.ModuleList([ |
| ReciprocalLayerwithCNN(d_model, d_inner, d_hidden, n_head, d_k, d_v) |
| for _ in range(n_layers)]) |
|
|
| self.dropout = nn.Dropout(dropout) |
| self.dropout_2 = nn.Dropout(dropout) |
|
|
| def forward(self, peptide_sequence, protein_sequence): |
| sequence_attention_list = [] |
|
|
| prot_attention_list = [] |
|
|
| prot_seq_attention_list = [] |
|
|
| seq_prot_attention_list = [] |
|
|
| sequence_enc = self.dropout(self.linear1(peptide_sequence)) |
|
|
| prot_enc = self.dropout_2(self.linear2(protein_sequence)) |
|
|
| for reciprocal_layer in self.reciprocal_layer_stack: |
| prot_enc, sequence_enc, prot_attention, sequence_attention, prot_seq_attention, seq_prot_attention = \ |
| reciprocal_layer(sequence_enc, prot_enc) |
|
|
| sequence_attention_list.append(sequence_attention) |
|
|
| prot_attention_list.append(prot_attention) |
|
|
| prot_seq_attention_list.append(prot_seq_attention) |
|
|
| seq_prot_attention_list.append(seq_prot_attention) |
|
|
| return prot_enc, sequence_enc, sequence_attention_list, prot_attention_list, \ |
| seq_prot_attention_list, seq_prot_attention_list |
|
|
|
|
| class RepeatedModule2(nn.Module): |
| def __init__(self, n_layers, d_model, |
| n_head, d_k, d_v, d_inner, dropout=0.1): |
| super().__init__() |
|
|
| self.linear1 = nn.Linear(1280, d_model) |
| self.linear2 = nn.Linear(1280, d_model) |
| self.sequence_embedding = nn.Embedding(20, d_model) |
| self.d_model = d_model |
|
|
| self.reciprocal_layer_stack = nn.ModuleList([ |
| ReciprocalLayer(d_model, d_inner, n_head, d_k, d_v) |
| for _ in range(n_layers)]) |
|
|
| self.dropout = nn.Dropout(dropout) |
| self.dropout_2 = nn.Dropout(dropout) |
|
|
| def forward(self, peptide_sequence, protein_sequence): |
| sequence_attention_list = [] |
|
|
| prot_attention_list = [] |
|
|
| prot_seq_attention_list = [] |
|
|
| seq_prot_attention_list = [] |
|
|
| sequence_enc = self.dropout(self.linear1(peptide_sequence)) |
|
|
| prot_enc = self.dropout_2(self.linear2(protein_sequence)) |
|
|
| for reciprocal_layer in self.reciprocal_layer_stack: |
| prot_enc, sequence_enc, prot_attention, sequence_attention, prot_seq_attention, seq_prot_attention = \ |
| reciprocal_layer(sequence_enc, prot_enc) |
|
|
| sequence_attention_list.append(sequence_attention) |
|
|
| prot_attention_list.append(prot_attention) |
|
|
| prot_seq_attention_list.append(prot_seq_attention) |
|
|
| seq_prot_attention_list.append(seq_prot_attention) |
|
|
| return prot_enc, sequence_enc, sequence_attention_list, prot_attention_list, \ |
| seq_prot_attention_list, seq_prot_attention_list |
|
|
|
|
| class RepeatedModule(nn.Module): |
| |
| def __init__(self, n_layers, d_model, |
| n_head, d_k, d_v, d_inner, dropout=0.1): |
| |
| super().__init__() |
| |
| self.linear = nn.Linear(1024, d_model) |
| self.sequence_embedding = nn.Embedding(20, d_model) |
| self.d_model = d_model |
| |
| self.reciprocal_layer_stack = nn.ModuleList([ |
| ReciprocalLayer(d_model, d_inner, n_head, d_k, d_v) |
| for _ in range(n_layers)]) |
| |
| self.dropout = nn.Dropout(dropout) |
| self.dropout_2 = nn.Dropout(dropout) |
|
|
| |
| |
| def _positional_embedding(self, batches, number): |
| |
| result = torch.exp(torch.arange(0, self.d_model,2,dtype=torch.float32)*-1*(np.log(10000)/self.d_model)) |
| |
| numbers = torch.arange(0, number, dtype=torch.float32) |
| |
| numbers = numbers.unsqueeze(0) |
| |
| numbers = numbers.unsqueeze(2) |
| |
| result = numbers*result |
| |
| result = torch.cat((torch.sin(result), torch.cos(result)),2) |
| |
| return result |
| |
| def forward(self, peptide_sequence, protein_sequence): |
| |
| |
| sequence_attention_list = [] |
| |
| prot_attention_list = [] |
| |
| prot_seq_attention_list = [] |
| |
| seq_prot_attention_list = [] |
| |
| sequence_enc = self.sequence_embedding(peptide_sequence) |
| |
| sequence_enc += to_var(self._positional_embedding(peptide_sequence.shape[0], |
| peptide_sequence.shape[1])) |
| sequence_enc = self.dropout(sequence_enc) |
|
|
|
|
|
|
|
|
|
|
| prot_enc = self.dropout_2(self.linear(protein_sequence)) |
| |
| |
| |
|
|
| for reciprocal_layer in self.reciprocal_layer_stack: |
| |
| prot_enc, sequence_enc, prot_attention, sequence_attention, prot_seq_attention, seq_prot_attention =\ |
| reciprocal_layer(sequence_enc, prot_enc) |
| |
| sequence_attention_list.append(sequence_attention) |
| |
| prot_attention_list.append(prot_attention) |
| |
| prot_seq_attention_list.append(prot_seq_attention) |
| |
| seq_prot_attention_list.append(seq_prot_attention) |
| |
| |
| |
| return prot_enc, sequence_enc, sequence_attention_list, prot_attention_list,\ |
| seq_prot_attention_list, seq_prot_attention_list |
|
|
|
|
| class FullModel(nn.Module): |
|
|
| def __init__(self, n_layers, d_model, n_head, |
| d_k, d_v, d_inner, return_attention=False, dropout=0.2): |
| super().__init__() |
|
|
| self.esm_model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D") |
|
|
| |
| for param in self.esm_model.parameters(): |
| param.requires_grad = False |
|
|
| self.repeated_module = RepeatedModule2(n_layers, d_model, |
| 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.sigmoid = nn.Sigmoid() |
|
|
| self.return_attention = return_attention |
|
|
| def forward(self, binder_tokens, target_tokens): |
|
|
| with torch.no_grad(): |
| 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.sigmoid(self.output_projection_prot(prot_enc)) |
|
|
| return prot_enc |
|
|
|
|
|
|
| class Original_FullModel(nn.Module): |
| |
| def __init__(self, n_layers, d_model, n_head, |
| d_k, d_v, d_inner, return_attention=False, dropout=0.2): |
| |
| super().__init__() |
| self.repeated_module = RepeatedModule(n_layers, d_model, |
| 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, 2) |
| |
| |
| |
| self.softmax_prot =nn.LogSoftmax(dim=-1) |
| |
| |
| self.return_attention = return_attention |
| |
| def forward(self, peptide_sequence, protein_sequence): |
|
|
| 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.softmax_prot(self.output_projection_prot(prot_enc)) |
| |
| |
| |
| |
| |
| if not self.return_attention: |
| return prot_enc |
| else: |
| return prot_enc, sequence_attention_list, prot_attention_list,\ |
| seq_prot_attention_list, seq_prot_attention_list |
| |
|
|
