AptaBLE / encoders.py
Atom Bioworks
Create encoders.py
20c9cc4 verified
raw
history blame
10.7 kB
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