try protein sequence prediction directly with esm package

app.py

@@ -4,7 +4,7 @@ import torch
 import numpy as np 
 import pandas as pd
 import streamlit as st
-import bio_embeddings.embed
+import esm
 
 from rdkit import Chem
 from rdkit.Chem import Draw
@@ -191,8 +191,13 @@ def display_protein():
     sequence = st.text_input('Enter the amino-acid sequence of the query protein target', value='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA', placeholder='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA')
     
     if sequence:
-        embeddings = bio_embeddings.embed(sequence, "esm")
-        coordinates = embeddings["esm"].detach().numpy()
+        model, alphabet = esm.pretrained.esm1_t6_43M_UR50S()
+
+        batch_converter = alphabet.get_batch_converter()
+        inputs = batch_converter([sequence])
+        embeddings = model.embed(inputs)
+        coordinates = embeddings["representations"][0].detach().numpy()
+        
         st.write(coordinates)
     
     """