Added msa and go stubs

app.py

@@ -1,5 +1,5 @@
 # credit: https://huggingface.co/spaces/simonduerr/3dmol.js/blob/main/app.py
-
+from typing import Tuple
 import os
 import sys
 from urllib import request
@@ -9,6 +9,9 @@ import requests
 from transformers import AutoTokenizer, AutoModelForMaskedLM, EsmModel, AutoModel
 import torch
 import progres as pg
+import esm
+
+import msa
 
 
 tokenizer_nt = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")
@@ -23,6 +26,11 @@ tokenizer_se = AutoTokenizer.from_pretrained('sentence-transformers/all-mpnet-ba
 model_se = AutoModel.from_pretrained('sentence-transformers/all-mpnet-base-v2')
 model_se.eval()
 
+msa_transformer, msa_transformer_alphabet = esm.pretrained.esm_msa1b_t12_100M_UR50S()
+msa_transformer = msa_transformer.eval()
+msa_transformer_batch_converter = msa_transformer_alphabet.get_batch_converter()
+
+
 
 def nt_embed(sequence: str):
     tokens_ids = tokenizer_nt.batch_encode_plus([sequence], return_tensors="pt")["input_ids"]
@@ -51,6 +59,17 @@ def se_embed(sentence: str):
     return model_output[0]
 
 
+def msa_embed(msa):
+    inputs = msa.greedy_select(inputs, 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)
+    
+    temp = msa_transformer(msa_transformer_batch_tokens,repr_layers=[12])['representations']
+    temp = temp[12][:,:,0,:]
+    temp = torch.mean(temp,(0,1))
+    return temp
+
+
 def download_data_if_required():
     url_base = f"https://zenodo.org/record/{pg.zenodo_record}/files"
     fps = [pg.trained_model_fp]
@@ -181,6 +200,14 @@ def update_se(inp):
     return str(se_embed(inp))
 
 
+def update_go(inp):
+    return str(go_embed(inp))
+
+
+def update_msa(inp):
+    return str(msa_embed(msa.read_msa(inp)))
+
+
 demo = gr.Blocks()
 
 with demo:
@@ -222,6 +249,21 @@ with demo:
                 btn = gr.Button("View embeddings")
                 emb = gr.Textbox(interactive=False)
                 btn.click(fn=update_se, inputs=[inp], outputs=emb)
+        with gr.TabItem("MSA Embeddings"):
+            with gr.Box():
+                inp = gr.File(file_count="single", label="Input MSA")
+                btn = gr.Button("View embeddings")
+                emb = gr.Textbox(interactive=False)
+                btn.click(fn=update_msa, inputs=[inp], outputs=emb)
+        with gr.TabItem("GO Embeddings"):
+            with gr.Box():
+                inp = gr.Textbox(
+                    placeholder="", label="Input GO Terms"
+                )
+                btn = gr.Button("View embeddings")
+                emb = gr.Textbox(interactive=False)
+                btn.click(fn=update_go, inputs=[inp], outputs=emb)
+
 
 if __name__ == "__main__":
     download_data_if_required()

requirements.txt

@@ -5,8 +5,9 @@ requests==2.31.0
 torch==2.0.1
 --find-links https://data.pyg.org/whl/torch-2.0.0+cpu.html torch-cluster==1.6.1
 torch-geometric==2.3.1
---find-links https://data.pyg.org/whl/torch-2.0.0+cpu.html torch-scatter==2.1.1
+torch-scatter==2.1.1
 --find-links https://data.pyg.org/whl/torch-2.0.0+cpu.html torch-sparse==0.6.17
 --find-links https://data.pyg.org/whl/torch-2.0.0+cpu.html torch-spline-conv==1.2.2
 transformers==4.29.2
-progres
+progres
+fair-esm