feat: Add get_esm2_embeddings function

app.py

@@ -38,6 +38,22 @@ def fold_prot_locally(sequence):
     pdb = convert_outputs_to_pdb(output)
     return pdb
 
+def get_esm2_embeddings(sequence):
+    logger.info("Getting embeddings for: " + sequence)
+    tokenized_input = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids'].cuda()
+
+    with torch.no_grad():
+        aa = tokenized_input 
+        L = aa.shape[1]
+        device = tokenized_input.device
+        attention_mask = torch.ones_like(aa, device=device)
+
+        # === ESM ===
+        esmaa = model.af2_idx_to_esm_idx(aa, attention_mask)
+        esm_s = model.compute_language_model_representations(esmaa)
+
+    return {"res": esm_s.cpu().tolist()}
+
 def get_esmfold_embeddings(sequence):
     logger.info("Getting embeddings for: " + sequence)
     tokenized_input = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids'].cuda()
@@ -165,11 +181,12 @@ with gr.Blocks() as demo:
     with gr.Row(visible=False):
        with gr.Column():
            gr.Markdown("## Embeddings")
-           embs = gr.JSON(label="Embeddings", interactive=False)
+           embs = gr.JSON(label="Embeddings")
 
     name.change(fn=suggest, inputs=name, outputs=inp)
     btn.click(fold_prot_locally, inputs=[inp], outputs=[out], api_name="pdb")
     btn.click(get_esmfold_embeddings, inputs=[inp], outputs=[embs], api_name="embeddings")
+    btn.click(get_esm2_embeddings, inputs=[inp], outputs=[embs], api_name="esm2_embeddings")
     out.change(fn=molecule, inputs=[out], outputs=[out_mol], api_name="3d_fold")
 
 demo.launch()

client.py

@@ -0,0 +1,14 @@
+#%%
+from gradio_client import Client
+
+#%%
+# client = Client("https://wwydmanski-esmfold.hf.space/")
+client = Client("http://localhost:7860")
+
+# %%
+result = client.predict(
+    "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGGGPGAGSLQPLALEGSLQKRGIVEQCCTSICSLYQLENYCN",	# str in 'sequence' Textbox component
+    api_name="/esm2_embeddings")
+
+# %%
+result