Update app.py

app.py

@@ -9,15 +9,21 @@ import spaces
 import gradio as gr
 from esm_scripts.extract import run_demo
 from esm import pretrained, FastaBatchedDataset
+from transformers import EsmTokenizer, EsmModel
+
 
 # Load the model
 model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
 model.load_checkpoint("model/checkpoint_mf2.pth")
 model.to('cuda')
 
-model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
-model_esm.to('cuda')
-model_esm.eval()
+# model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
+# model_esm.to('cuda')
+# model_esm.eval()
+tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t36_3B_UR50D")
+model = EsmModel.from_pretrained("facebook/esm2_t36_3B_UR50D")
+model.to('cuda')
+model.eval()
 
 @spaces.GPU
 def generate_caption(protein, prompt):
@@ -29,7 +35,8 @@ def generate_caption(protein, prompt):
     # esm_emb = run_demo(protein_name='protein_name', protein_seq=protein, 
     #                    model=model_esm, alphabet=alphabet, 
     #                    include='per_tok', repr_layers=[36], truncation_seq_length=1024)
-    protein_name='protein_name'
+    '''
+        protein_name='protein_name'
     protein_seq=protein
     include='per_tok'
     repr_layers=[36]
@@ -86,6 +93,11 @@ def generate_caption(protein, prompt):
                 if return_contacts:
                     result["contacts"] = contacts[i, : truncate_len, : truncate_len].clone()
             esm_emb = result['representations'][36]
+    '''
+    inputs = tokenizer([protein], return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    esm_emb = outputs.last_hidden_state.detach()[0]
     
     print("esm embedding generated")
     esm_emb = F.pad(esm_emb.t(), (0, 1024 - len(esm_emb))).t().to('cuda')