Update app.py

app.py

@@ -9,7 +9,7 @@ import spaces
 import gradio as gr
 from esm_scripts.extract import run_demo
 from esm import pretrained, FastaBatchedDataset
-from transformers import EsmTokenizer, EsmModel
+# from transformers import EsmTokenizer, EsmModel
 
 
 # Load the model
@@ -17,13 +17,13 @@ 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()
-tokenizer = EsmTokenizer.from_pretrained("facebook/esm2_t36_3B_UR50D")
-model_esm = EsmModel.from_pretrained("facebook/esm2_t36_3B_UR50D")
+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_esm = EsmModel.from_pretrained("facebook/esm2_t36_3B_UR50D")
+# model_esm.to('cuda')
+# model_esm.eval()
 
 @spaces.GPU
 def generate_caption(protein, prompt):
@@ -35,8 +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]
@@ -98,7 +98,7 @@ def generate_caption(protein, prompt):
     with torch.no_grad():
         outputs = model_esm(**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')
     print("esm embedding processed")
@@ -117,6 +117,7 @@ description = """Quick demonstration of the FAPM model for protein function pred
 
 The model used in this app is available at [Hugging Face Model Hub](https://huggingface.co/wenkai/FAPM) and the source code can be found on [GitHub](https://github.com/xiangwenkai/FAPM/tree/main)."""
 
+
 iface = gr.Interface(
     fn=generate_caption,
     inputs=[gr.Textbox(type="text", label="Upload sequence"), gr.Textbox(type="text", label="Prompt")],