| import torch |
| import torch.nn as nn |
| from typing import List, Dict |
|
|
| class AlphaFoldEmbedTokenizer: |
| def __init__(self): |
| self.vocab = {aa: i for i, aa in enumerate("ACDEFGHIKLMNPQRSTVWY")} |
| self.vocab["<UNK>"] = len(self.vocab) |
| |
| def tokenize(self, seq: str) -> List[int]: |
| return [self.vocab.get(aa, self.vocab["<UNK>"]) for aa in seq] |
|
|
| class AlphaFoldEmbedInference(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.embed = nn.Embedding(25, 512) |
| self.transformer = nn.TransformerEncoder( |
| nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True), |
| num_layers=4 |
| ) |
| self.plddt_head = nn.Linear(512, 1) |
| self.domain_head = nn.Linear(512, 1) |
|
|
| def forward(self, input_ids): |
| x = self.embed(input_ids) |
| x = self.transformer(x) |
| plddt = torch.sigmoid(self.plddt_head(x)).squeeze(-1) * 100 |
| domain_boundaries = torch.sigmoid(self.domain_head(x)).squeeze(-1) |
| return {"embeddings": x, "plddt": plddt, "domain_boundaries": domain_boundaries} |
|
|
| def process_proteins(proteins: List[str]): |
| tokenizer = AlphaFoldEmbedTokenizer() |
| model = AlphaFoldEmbedInference() |
| model.eval() |
| |
| results = [] |
| with torch.no_grad(): |
| for seq in proteins: |
| tokens = tokenizer.tokenize(seq) |
| input_ids = torch.tensor([tokens]) |
| out = model(input_ids) |
| results.append(out) |
| return results |
|
|
| if __name__ == "__main__": |
| res = process_proteins(["MKALIVLGLVLLSVTVQG", "MTEYKLVVVGAGGVGKSALTIQLIQNHFVDEYDPTIEDSYRKQVVIDGETCLLDILDTAGQ"]) |
| print("Done processing proteins.") |
|
|
