encoders.py

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