Added protein_bind methods

.gitignore

@@ -0,0 +1 @@
+env/

app.py

@@ -1,18 +1,18 @@
 # credit: https://huggingface.co/spaces/simonduerr/3dmol.js/blob/main/app.py
-from typing import Tuple
 import os
 import sys
 from urllib import request
 
+import esm
 import gradio as gr
+import progres as pg
 import requests
-from transformers import AutoTokenizer, AutoModelForMaskedLM, EsmModel, AutoModel
 import torch
-import progres as pg
-import esm
+from transformers import (AutoModel, AutoModelForMaskedLM, AutoTokenizer,
+                          EsmModel)
 
 import msa
-
+import proteinbind_new
 
 tokenizer_nt = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")
 model_nt = AutoModelForMaskedLM.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")
@@ -30,6 +30,15 @@ msa_transformer, msa_transformer_alphabet = esm.pretrained.esm_msa1b_t12_100M_UR
 msa_transformer = msa_transformer.eval()
 msa_transformer_batch_converter = msa_transformer_alphabet.get_batch_converter()
 
+model = proteinbind_new.create_proteinbind(True)
+
+
+def pass_through(torch_output, key: str):
+    input_data = {
+        key: torch_output,
+    }
+    output = model(input_data)
+    return output[key]
 
 
 def nt_embed(sequence: str):
@@ -37,38 +46,38 @@ def nt_embed(sequence: str):
     attention_mask = tokens_ids != tokenizer_nt.pad_token_id
     with torch.no_grad():
         torch_outs = model_nt(
-            tokens_ids,#.to('cuda'),
-            attention_mask=attention_mask,#.to('cuda'),
+            tokens_ids,  # .to('cuda'),
+            attention_mask=attention_mask,  # .to('cuda'),
             output_hidden_states=True
         )
     last_layer_CLS = torch_outs.hidden_states[-1].detach()[:, 0, :][0]
-    return last_layer_CLS
+    return pass_through(last_layer_CLS, "dna")
 
 
 def aa_embed(sequence: str):
     tokens = tokenizer_aa([sequence], return_tensors="pt")
     with torch.no_grad():
         torch_outs = model_aa(**tokens)
-    return torch_outs[0]
+    return pass_through(torch_outs[0], "aa")
 
 
 def se_embed(sentence: str):
     encoded_input = tokenizer_se([sentence], return_tensors='pt')
     with torch.no_grad():
         model_output = model_se(**encoded_input)
-    return model_output[0]
+    return pass_through(model_output[0], "text")
 
 
 def msa_embed(sequences: list):
-    inputs = msa.greedy_select(sequences, num_seqs=128) # can change this to pass more/fewer sequences
+    inputs = msa.greedy_select(sequences, num_seqs=128)  # can change this to pass more/fewer sequences
     msa_transformer_batch_labels, msa_transformer_batch_strs, msa_transformer_batch_tokens = msa_transformer_batch_converter([inputs])
     msa_transformer_batch_tokens = msa_transformer_batch_tokens.to(next(msa_transformer.parameters()).device)
-    
+
     with torch.no_grad():
-        temp = msa_transformer(msa_transformer_batch_tokens,repr_layers=[12])['representations']
-    temp = temp[12][:,:,0,:]
-    temp = torch.mean(temp,(0,1))
-    return temp
+        temp = msa_transformer(msa_transformer_batch_tokens, repr_layers=[12])['representations']
+    temp = temp[12][:, :, 0, :]
+    temp = torch.mean(temp, (0, 1))
+    return pass_through(temp, "msa")
 
 
 def go_embed(terms):
@@ -79,13 +88,13 @@ def download_data_if_required():
     url_base = f"https://zenodo.org/record/{pg.zenodo_record}/files"
     fps = [pg.trained_model_fp]
     urls = [f"{url_base}/trained_model.pt"]
-    #for targetdb in pre_embedded_dbs:
+    # for targetdb in pre_embedded_dbs:
     #    fps.append(os.path.join(database_dir, targetdb + ".pt"))
     #    urls.append(f"{url_base}/{targetdb}.pt")
 
     if not os.path.isdir(pg.trained_model_dir):
         os.makedirs(pg.trained_model_dir)
-    #if not os.path.isdir(database_dir):
+    # if not os.path.isdir(database_dir):
     #    os.makedirs(database_dir)
 
     printed = False
@@ -103,7 +112,7 @@ def download_data_if_required():
                     assert "model" in d
                 else:
                     assert "embeddings" in d
-            except:
+            except Exception:
                 if os.path.isfile(fp):
                     os.remove(fp)
                 print("Failed to download from", url, "and save to", fp, file=sys.stderr)
@@ -119,7 +128,7 @@ def get_pdb(pdb_code="", filepath=""):
         try:
             with open(filepath.name) as f:
                 return f.read()
-        except AttributeError as e:
+        except AttributeError:
             return None
     else:
         return requests.get(f"https://files.rcsb.org/view/{pdb_code}.pdb").content.decode()
@@ -150,12 +159,12 @@ def molecule(pdb):
     </head>
     <body>  
     <div id="container" class="mol-container"></div>
-  
+
             <script>
                let pdb = `"""
         + pdb
         + """`  
-      
+
              $(document).ready(function () {
                 let element = $("#container");
                 let config = { backgroundColor: "black" };
@@ -272,4 +281,4 @@ with demo:
 
 if __name__ == "__main__":
     download_data_if_required()
-    demo.launch()
+    demo.launch()

proteinbind_new.py

@@ -15,6 +15,7 @@ ModalityType = SimpleNamespace(
     TEXT="text",
 )
 
+
 class Normalize(nn.Module):
     def __init__(self, dim: int) -> None:
         super().__init__()
@@ -23,6 +24,7 @@ class Normalize(nn.Module):
     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
@@ -42,6 +44,7 @@ class EmbeddingDataset(Dataset):
         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
@@ -60,7 +63,8 @@ class DualEmbeddingDataset(Dataset):
         sequence_embedding = self.sequence_embedding[idx]
         embedding = self.embedding[idx]
         return {"aa": sequence_embedding, self.modality: embedding}
-    
+
+
 class ProteinBindModel(nn.Module):
 
     def __init__(
@@ -92,7 +96,6 @@ class ProteinBindModel(nn.Module):
             out_embed_dim
         )
 
-
     def _create_modality_trunk(
             self,
             aa_embed_dim,
@@ -140,7 +143,7 @@ class ProteinBindModel(nn.Module):
             nn.ReLU(),
             nn.Linear(512, in_embed_dim),
         )
-        
+
         modality_trunks[ModalityType.GO] = nn.Sequential(
             nn.Linear(go_embed_dim, 512),
             nn.ReLU(),
@@ -220,7 +223,6 @@ class ProteinBindModel(nn.Module):
         modality_postprocessors[ModalityType.GO] = Normalize(dim=-1)
         modality_postprocessors[ModalityType.MSA] = Normalize(dim=-1)
 
-
         return nn.ModuleDict(modality_postprocessors)
 
     def forward(self, inputs):
@@ -239,7 +241,6 @@ class ProteinBindModel(nn.Module):
 
         for modality_key, modality_value in inputs.items():
 
-
             modality_value = self.modality_trunks[modality_key](
                 modality_value
             )
@@ -247,10 +248,10 @@ class ProteinBindModel(nn.Module):
             modality_value = self.modality_heads[modality_key](
                 modality_value
             )
-             
+
             modality_value = self.modality_postprocessors[modality_key](
-                    modality_value
-                )
+                modality_value
+            )
             outputs[modality_key] = modality_value
 
         return outputs
@@ -274,7 +275,7 @@ def create_proteinbind(pretrained=False):
     )
 
     if pretrained:
-        #get path from config
+        # get path from config
         PATH = 'best_model.pth'
 
         model.load_state_dict(torch.load(PATH))