feat: add embeddings

app.py

@@ -38,6 +38,15 @@ def fold_prot_locally(sequence):
     pdb = convert_outputs_to_pdb(output)
     return pdb
 
+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()
+
+    with torch.no_grad():
+        output = model(tokenized_input)
+
+    return output["s_s"].cpu().numpy()
+
 def suggest(option):
    if option == "Plastic degradation protein":
      suggestion = "MGSSHHHHHHSSGLVPRGSHMRGPNPTAASLEASAGPFTVRSFTVSRPSGYGAGTVYYPTNAGGTVGAIAIVPGYTARQSSIKWWGPRLASHGFVVITIDTNSTLDQPSSRSSQQMAALRQVASLNGTSSSPIYGKVDTARMGVMGWSMGGGGSLISAANNPSLKAAAPQAPWDSSTNFSSVTVPTLIFACENDSIAPVNSSALPIYDSMSRNAKQFLEINGGSHSCANSGNSNQALIGKKGVAWMKRFMDNDTRYSTFACENPNSTRVSDFRTANCSLEDPAANKARKEAELAAATAEQ"
@@ -153,8 +162,14 @@ with gr.Blocks() as demo:
         with gr.Column():
             out_mol = gr.HTML(label="3D Structure")
 
+    with gr.Row(visible=False):
+       with gr.Column():
+           gr.Markdown("## Embeddings")
+           embs = gr.Code(label="Embeddings", interactive=False)
+
     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")
     out.change(fn=molecule, inputs=[out], outputs=[out_mol], api_name="3d_fold")
 
 demo.launch()