|
from types import SimpleNamespace |
|
|
|
import pandas as pd |
|
import torch |
|
import torch.nn as nn |
|
from torch.utils.data import Dataset |
|
|
|
|
|
ModalityType = SimpleNamespace( |
|
AA="aa", |
|
DNA="dna", |
|
PDB="pdb", |
|
GO="go", |
|
MSA="msa", |
|
TEXT="text", |
|
) |
|
|
|
class Normalize(nn.Module): |
|
def __init__(self, dim: int) -> None: |
|
super().__init__() |
|
self.dim = dim |
|
|
|
def forward(self, x): |
|
return torch.nn.functional.normalize(x, dim=self.dim, p=2) |
|
|
|
class EmbeddingDataset(Dataset): |
|
""" |
|
The main class for turning any modality to a torch Dataset that can be passed to |
|
a torch dataloader. Any modality that doesn't fit into the __getitem__ |
|
method can subclass this and modify the __getitem__ method. |
|
""" |
|
def __init__(self, sequence_file_path, embeddings_file_path, modality): |
|
self.sequence = pd.read_csv(sequence_file_path) |
|
self.embedding = torch.load(embeddings_file_path) |
|
self.modality = modality |
|
|
|
def __len__(self): |
|
return len(self.sequence) |
|
|
|
def __getitem__(self, idx): |
|
sequence = self.sequence.iloc[idx, 0] |
|
embedding = self.embedding[idx] |
|
return {"aa": sequence, self.modality: embedding} |
|
|
|
class DualEmbeddingDataset(Dataset): |
|
""" |
|
The main class for turning any modality to a torch Dataset that can be passed to |
|
a torch dataloader. Any modality that doesn't fit into the __getitem__ |
|
method can subclass this and modify the __getitem__ method. |
|
""" |
|
def __init__(self, sequence_embeddings_file_path, embeddings_file_path, modality): |
|
self.sequence_embedding = torch.load(sequence_embeddings_file_path) |
|
self.embedding = torch.load(embeddings_file_path) |
|
self.modality = modality |
|
|
|
def __len__(self): |
|
return len(self.sequence_embedding) |
|
|
|
def __getitem__(self, idx): |
|
sequence_embedding = self.sequence_embedding[idx] |
|
embedding = self.embedding[idx] |
|
return {"aa": sequence_embedding, self.modality: embedding} |
|
|
|
class ProteinBindModel(nn.Module): |
|
|
|
def __init__( |
|
self, |
|
aa_embed_dim, |
|
dna_embed_dim, |
|
pdb_embed_dim, |
|
go_embed_dim, |
|
msa_embed_dim, |
|
text_embed_dim, |
|
in_embed_dim, |
|
out_embed_dim |
|
): |
|
super().__init__() |
|
self.modality_trunks = self._create_modality_trunk( |
|
aa_embed_dim, |
|
dna_embed_dim, |
|
pdb_embed_dim, |
|
go_embed_dim, |
|
msa_embed_dim, |
|
text_embed_dim, |
|
out_embed_dim |
|
) |
|
self.modality_heads = self._create_modality_head( |
|
in_embed_dim, |
|
out_embed_dim, |
|
) |
|
self.modality_postprocessors = self._create_modality_postprocessors( |
|
out_embed_dim |
|
) |
|
|
|
|
|
def _create_modality_trunk( |
|
self, |
|
aa_embed_dim, |
|
dna_embed_dim, |
|
pdb_embed_dim, |
|
go_embed_dim, |
|
msa_embed_dim, |
|
text_embed_dim, |
|
in_embed_dim |
|
): |
|
""" |
|
The current layers are just a proof of concept |
|
and are subject to the opinion of others. |
|
:param aa_embed_dim: |
|
:param dna_embed_dim: |
|
:param pdb_embed_dim: |
|
:param go_embed_dim: |
|
:param msa_embed_dim: |
|
:param text_embed_dim: |
|
:param in_embed_dim: |
|
:return: |
|
""" |
|
modality_trunks = {} |
|
|
|
modality_trunks[ModalityType.AA] = nn.Sequential( |
|
nn.Linear(aa_embed_dim, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, in_embed_dim), |
|
) |
|
|
|
modality_trunks[ModalityType.DNA] = nn.Sequential( |
|
nn.Linear(dna_embed_dim, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, in_embed_dim), |
|
) |
|
|
|
modality_trunks[ModalityType.PDB] = nn.Sequential( |
|
nn.Linear(pdb_embed_dim, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, in_embed_dim), |
|
) |
|
|
|
modality_trunks[ModalityType.GO] = nn.Sequential( |
|
nn.Linear(go_embed_dim, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, in_embed_dim), |
|
) |
|
|
|
modality_trunks[ModalityType.MSA] = nn.Sequential( |
|
nn.Linear(msa_embed_dim, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, in_embed_dim), |
|
) |
|
|
|
modality_trunks[ModalityType.TEXT] = nn.Sequential( |
|
nn.Linear(text_embed_dim, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, 512), |
|
nn.ReLU(), |
|
nn.Linear(512, in_embed_dim), |
|
) |
|
|
|
return nn.ModuleDict(modality_trunks) |
|
|
|
def _create_modality_head( |
|
self, |
|
in_embed_dim, |
|
out_embed_dim |
|
): |
|
modality_heads = {} |
|
|
|
modality_heads[ModalityType.AA] = nn.Sequential( |
|
nn.LayerNorm(normalized_shape=in_embed_dim, eps=1e-6), |
|
nn.Dropout(p=0.5), |
|
nn.Linear(in_embed_dim, out_embed_dim, bias=False), |
|
) |
|
|
|
modality_heads[ModalityType.DNA] = nn.Sequential( |
|
nn.LayerNorm(normalized_shape=in_embed_dim, eps=1e-6), |
|
nn.Dropout(p=0.5), |
|
nn.Linear(in_embed_dim, out_embed_dim, bias=False), |
|
) |
|
|
|
modality_heads[ModalityType.PDB] = nn.Sequential( |
|
nn.LayerNorm(normalized_shape=in_embed_dim, eps=1e-6), |
|
nn.Dropout(p=0.5), |
|
nn.Linear(in_embed_dim, out_embed_dim, bias=False), |
|
) |
|
|
|
modality_heads[ModalityType.GO] = nn.Sequential( |
|
nn.LayerNorm(normalized_shape=in_embed_dim, eps=1e-6), |
|
nn.Dropout(p=0.5), |
|
nn.Linear(in_embed_dim, out_embed_dim, bias=False), |
|
) |
|
|
|
modality_heads[ModalityType.MSA] = nn.Sequential( |
|
nn.LayerNorm(normalized_shape=in_embed_dim, eps=1e-6), |
|
nn.Dropout(p=0.5), |
|
nn.Linear(in_embed_dim, out_embed_dim, bias=False), |
|
) |
|
|
|
modality_heads[ModalityType.TEXT] = nn.Sequential( |
|
nn.LayerNorm(normalized_shape=in_embed_dim, eps=1e-6), |
|
nn.Dropout(p=0.5), |
|
nn.Linear(in_embed_dim, out_embed_dim, bias=False), |
|
) |
|
return nn.ModuleDict(modality_heads) |
|
|
|
def _create_modality_postprocessors(self, out_embed_dim): |
|
modality_postprocessors = {} |
|
modality_postprocessors[ModalityType.AA] = Normalize(dim=-1) |
|
modality_postprocessors[ModalityType.DNA] = Normalize(dim=-1) |
|
modality_postprocessors[ModalityType.PDB] = Normalize(dim=-1) |
|
modality_postprocessors[ModalityType.TEXT] = Normalize(dim=-1) |
|
modality_postprocessors[ModalityType.GO] = Normalize(dim=-1) |
|
modality_postprocessors[ModalityType.MSA] = Normalize(dim=-1) |
|
|
|
|
|
return nn.ModuleDict(modality_postprocessors) |
|
|
|
def forward(self, inputs): |
|
""" |
|
input = {k_1: [v],k_n: [v]} |
|
for key in input |
|
get trunk for key |
|
forward pass of value in trunk |
|
get projection head of key |
|
forward pass of value in projection head |
|
append output in output dict |
|
return { k_1, [o], k_n: [o]} |
|
""" |
|
|
|
outputs = {} |
|
|
|
for modality_key, modality_value in inputs.items(): |
|
|
|
|
|
modality_value = self.modality_trunks[modality_key]( |
|
modality_value |
|
) |
|
|
|
modality_value = self.modality_heads[modality_key]( |
|
modality_value |
|
) |
|
|
|
modality_value = self.modality_postprocessors[modality_key]( |
|
modality_value |
|
) |
|
outputs[modality_key] = modality_value |
|
|
|
return outputs |
|
|
|
|
|
def create_proteinbind(pretrained=False): |
|
""" |
|
The embedding dimensions here are dummy |
|
:param pretrained: |
|
:return: |
|
""" |
|
model = ProteinBindModel( |
|
aa_embed_dim=480, |
|
dna_embed_dim=1280, |
|
pdb_embed_dim=128, |
|
go_embed_dim=600, |
|
msa_embed_dim=768, |
|
text_embed_dim=768, |
|
in_embed_dim=1024, |
|
out_embed_dim=1024 |
|
) |
|
|
|
if pretrained: |
|
|
|
PATH = 'best_model.pth' |
|
|
|
model.load_state_dict(torch.load(PATH)) |
|
|
|
return model |
|
|
