Create encoders.py

encoders.py

@@ -0,0 +1,264 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.weight_norm import weight_norm
+import math
+import numpy as np
+
+class cross_attn_block(nn.Module):
+    def __init__(self, embed_dim, n_heads, dropout):
+        super().__init__()
+        self.heads = n_heads
+        self.mha = nn.MultiheadAttention(embed_dim, n_heads, dropout, batch_first=True)
+        self.ln_apt = nn.LayerNorm(embed_dim)
+        self.ln_prot = nn.LayerNorm(embed_dim)
+        self.ln_out = nn.LayerNorm(embed_dim)
+        self.linear = nn.Linear(embed_dim, embed_dim)
+    
+    def forward(self, embeddings_x, embeddings_y, x_t, y_t):
+
+        # compute attention masks
+        attn_mask = generate_3d_mask(y_t, x_t, self.heads)
+
+        # apply layer norms
+        embeddings_x_n = self.ln_apt(embeddings_x)
+        embeddings_y_n = self.ln_prot(embeddings_y)
+
+        # perform cross-attention
+        reps = embeddings_y + self.mha(embeddings_y_n, embeddings_x_n, embeddings_x_n, attn_mask=attn_mask)[0]
+        return reps + self.linear(self.ln_out(reps))
+
+class self_attn_block(nn.Module):
+    def __init__(self, d_embed, heads, dropout):
+        super().__init__()
+        # self.l1 = nn.Linear(d_linear, d_linear)
+        self.heads = heads
+        self.ln1 = nn.LayerNorm(d_embed)
+        self.ln2 = nn.LayerNorm(d_embed)
+        self.mha = nn.MultiheadAttention(d_embed, self.heads, dropout, batch_first=True)
+        self.linear = nn.Linear(d_embed, d_embed)
+    
+    def forward(self, embeddings_x, x_t):
+
+        # compute attention masks
+        # attn_mask = generate_3d_mask(x_t, x_t, self.heads)
+        # apply layer norm
+        embeddings_x_n = self.ln1(embeddings_x)
+        reps = embeddings_x + self.mha(embeddings_x_n, embeddings_x_n, embeddings_x_n, key_padding_mask=~x_t)[0]
+        return reps + self.linear(self.ln2(reps))
+
+
+class AptaBLE(nn.Module):
+    def __init__(self, apta_encoder, prot_encoder, dropout):
+        super(AptaBLE, self).__init__()
+
+        #hyperparameters
+        self.apta_encoder = apta_encoder
+        self.prot_encoder = prot_encoder
+
+        self.flatten = nn.Flatten()
+        self.prot_reshape = nn.Linear(1280, 512)
+        self.apta_keep = nn.Linear(512, 512)
+
+        self.l1 = nn.Linear(1024, 1024)
+        self.l2 = nn.Linear(1024, 512)
+        self.l3 = nn.Linear(512, 256)
+        self.l4 = nn.Linear(256, 1)
+        self.can = CAN(512, 8, 1, 'mean_all_tok')
+        self.bn1 = nn.BatchNorm1d(1024)
+        self.bn2 = nn.BatchNorm1d(512)
+        self.bn3 = nn.BatchNorm1d(256)
+        self.relu = nn.ReLU()
+
+
+
+    def forward(self, apta_in, esm_prot, apta_attn, prot_attn):
+        apta = self.apta_encoder(apta_in, apta_attn, apta_attn, output_hidden_states=True)['hidden_states'][-1] # output: (BS X #apt_toks x apt_embed_dim), encoder outputs (BS x MLM & sec. structure feature embeddings)
+        
+        prot = self.prot_encoder(esm_prot, repr_layers=[33], return_contacts=False)['representations'][33]
+
+        prot = self.prot_reshape(prot)
+        apta = self.apta_keep(apta)
+
+        output, cross_map, prot_map, apta_map = self.can(prot, apta, prot_attn, apta_attn)
+        output = self.relu(self.l1(output))
+        output = self.bn1(output)
+        output = self.relu(self.l2(output))
+        output = self.bn2(output)
+        output = self.relu(self.l3(output))
+        output = self.bn3(output)
+        output = self.l4(output)
+        output = torch.sigmoid(output)
+
+        return output, cross_map, prot_map, apta_map
+
+def find_opt_threshold(target, pred):
+    result = 0
+    best = 0
+
+    for i in range(0, 1000):
+        pred_threshold = np.where(pred > i/1000, 1, 0)
+        now = f1_score(target, pred_threshold)
+        if now > best:
+            result = i/1000
+            best = now
+
+    return result
+
+def argument_seqset(seqset):
+    arg_seqset = []
+    for s, ss in seqset:
+        arg_seqset.append([s, ss]) 
+        
+        arg_seqset.append([s[::-1], ss[::-1]])
+
+    return arg_seqset
+
+def augment_apis(apta, prot, ys):
+    aug_apta = []
+    aug_prot = []
+    aug_y = []
+    for a, p, y in zip(apta, prot, ys):
+        aug_apta.append(a) 
+        aug_prot.append(p)
+        aug_y.append(y)
+        
+        aug_apta.append(a[::-1]) 
+        aug_prot.append(p)
+        aug_y.append(y)
+        
+        aug_apta.append(a) 
+        aug_prot.append(p[::-1])
+        aug_y.append(y)
+        
+        aug_apta.append(a[::-1]) 
+        aug_prot.append(p[::-1])
+        aug_y.append(y)
+        
+    return np.array(aug_apta), np.array(aug_prot), np.array(aug_y)
+
+def generate_3d_mask(batch1, batch2, heads):
+    # Ensure the batches are tensors
+    batch1 = torch.tensor(batch1, dtype=torch.bool)
+    batch2 = torch.tensor(batch2, dtype=torch.bool)
+    
+    # Validate that the batches have the same length
+    if batch1.size(0) != batch2.size(0):
+        raise ValueError("The batches must have the same number of vectors")
+    
+    # Generate the 3D mask for each pair of vectors
+    out_mask = []
+    masks = torch.stack([torch.ger(vec1, vec2) for vec1, vec2 in zip(batch1, batch2)])
+    for j in range(masks.shape[0]):
+        out_mask.append(torch.stack([masks[j] for i in range(heads)]))
+    # out_mask = torch.tensor(out_mask, dtype=bool)
+    out_mask = torch.cat(out_mask)
+    
+    # Replace False with -inf and True with 0
+    out_mask = out_mask.float()  # Convert to float to allow -inf
+    out_mask[out_mask == 0] = -1e9
+    out_mask[out_mask == 1] = 0
+    
+    return out_mask
+
+class CAN(nn.Module):
+    def __init__(self, hidden_dim, num_heads, group_size, aggregation):
+        super(CAN, self).__init__()
+        self.aggregation = aggregation
+        self.group_size = group_size
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.head_dim = hidden_dim // num_heads
+
+        # Protein weights
+        self.prot_query = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.prot_key = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.prot_val = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+        # Aptamer weights
+        self.apta_query = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.apta_key = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.apta_val = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+        # linear
+        self.lp = nn.Linear(hidden_dim, hidden_dim)
+
+    def mask_logits(self, logits, mask_row, mask_col, inf=1e6):
+        N, L1, L2, H = logits.shape
+        mask_row = mask_row.view(N, L1, 1).repeat(1, 1, H)
+        mask_col = mask_col.view(N, L2, 1).repeat(1, 1, H)
+
+        # Ignore all padding tokens across both embeddings
+        mask_pair = torch.einsum('blh, bkh->blkh', mask_row, mask_col)
+
+        # Set logit to -1e6 if masked
+        logits = torch.where(mask_pair, logits, logits - inf)
+        alpha = torch.softmax(logits, dim=2)
+        mask_row = mask_row.view(N, L1, 1, H).repeat(1, 1, L2, 1)
+        alpha = torch.where(mask_row, alpha, torch.zeros_like(alpha))
+        return alpha
+
+    def rearrange_heads(self, x, n_heads, n_ch):
+        # rearrange embedding for MHA
+        s = list(x.size())[:-1] + [n_heads, n_ch]
+        return x.view(*s)
+
+    def grouped_embeddings(self, x, mask, group_size):
+        N, L, D = x.shape
+        groups = L // group_size
+        # Average embeddings within each group
+        x_grouped = x.view(N, groups, group_size, D).mean(dim=2)
+        # Ignore groups without any non-padding tokens
+        mask_grouped = mask.view(N, groups, group_size).any(dim=2)
+        return x_grouped, mask_grouped
+
+    def forward(self, protein, aptamer, mask_prot, mask_apta):
+        # Group embeddings before applying multi-head attention
+        protein_grouped, mask_prot_grouped = self.grouped_embeddings(protein, mask_prot, self.group_size)
+        apta_grouped, mask_apta_grouped = self.grouped_embeddings(aptamer, mask_apta, self.group_size)
+
+        # Compute queries, keys, values for both protein and aptamer after grouping
+        query_prot = self.rearrange_heads(self.prot_query(protein_grouped), self.num_heads, self.head_dim)
+        key_prot = self.rearrange_heads(self.prot_key(protein_grouped), self.num_heads, self.head_dim)
+        value_prot = self.rearrange_heads(self.prot_val(protein_grouped), self.num_heads, self.head_dim)
+
+        query_apta = self.rearrange_heads(self.apta_query(apta_grouped), self.num_heads, self.head_dim)
+        key_apta = self.rearrange_heads(self.apta_key(apta_grouped), self.num_heads, self.head_dim)
+        value_apta = self.rearrange_heads(self.apta_val(apta_grouped), self.num_heads, self.head_dim)
+
+        # Compute attention scores
+        logits_pp = torch.einsum('blhd, bkhd->blkh', query_prot, key_prot)
+        logits_pa = torch.einsum('blhd, bkhd->blkh', query_prot, key_apta)
+        logits_ap = torch.einsum('blhd, bkhd->blkh', query_apta, key_prot)
+        logits_aa = torch.einsum('blhd, bkhd->blkh', query_apta, key_apta)
+
+        ml_pp = self.mask_logits(logits_pp, mask_prot_grouped, mask_prot_grouped)
+        ml_pa = self.mask_logits(logits_pa, mask_prot_grouped, mask_apta_grouped)
+        ml_ap = self.mask_logits(logits_ap, mask_apta_grouped, mask_prot_grouped)
+        ml_aa = self.mask_logits(logits_aa, mask_apta_grouped, mask_apta_grouped)
+
+        # Combine heads, combine self-attended and cross-attended representations (via avg)
+        prot_embedding = (torch.einsum('blkh, bkhd->blhd', ml_pp, value_prot).flatten(-2) +
+                   torch.einsum('blkh, bkhd->blhd', ml_pa, value_apta).flatten(-2)) / 2
+        apta_embedding = (torch.einsum('blkh, bkhd->blhd', ml_ap, value_prot).flatten(-2) +
+                   torch.einsum('blkh, bkhd->blhd', ml_aa, value_apta).flatten(-2)) / 2
+
+        prot_embedding += protein
+        apta_embedding += aptamer
+
+        # Aggregate token representations
+        if self.aggregation == "cls":
+            prot_embed = prot_embedding[:, 0]  # query : [batch_size, hidden]
+            apta_embed = apta_embedding[:, 0]  # query : [batch_size, hidden]
+        elif self.aggregation == "mean_all_tok":
+            prot_embed = prot_embedding.mean(1)  # query : [batch_size, hidden]
+            apta_embed = apta_embedding.mean(1)  # query : [batch_size, hidden]
+        elif self.aggregation == "mean":
+            prot_embed = (prot_embedding * mask_prot_grouped.unsqueeze(-1)).sum(1) / mask_prot_grouped.sum(-1).unsqueeze(-1)
+            apta_embed = (apta_embedding * mask_apta_grouped.unsqueeze(-1)).sum(1) / mask_apta_grouped.sum(-1).unsqueeze(-1)
+        else:
+            raise NotImplementedError()
+
+        embed = torch.cat([prot_embed, apta_embed], dim=1)
+
+        return embed, ml_pa, ml_pp, ml_aa