Update app.py

app.py

@@ -13,25 +13,35 @@ from data.evaluate_data.utils import Ontology
 import difflib
 import re
 from transformers import MistralForCausalLM
+from huggingface_hub import hf_hub_download
+bp_param = hf_hub_download(repo_id="wenkai/FAPM", filename="model/checkpoint_bp1.pth")
+mf_param = hf_hub_download(repo_id="wenkai/FAPM", filename="model/checkpoint_mf2.pth")
+cc_param = hf_hub_download(repo_id="wenkai/FAPM", filename="model/checkpoint_cc2.pth")
+# hf_hub_download(repo_id="wenkai/FAPM", filename="model/mf2_bert.pth")
+# hf_hub_download(repo_id="wenkai/FAPM", filename="model/bp1_bert.pth")
+# hf_hub_download(repo_id="wenkai/FAPM", filename="model/cc2_bert.pth")
 
+# bert_param = BertModel.from_pretrained("bert-base-uncased").state_dict()
 # Load the trained model
 def get_model(type='Molecule Function'):
     model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
     if type == 'Molecule Function':
-        model.load_checkpoint("model/checkpoint_mf2.pth")
+        # model.load_checkpoint("model/checkpoint_mf2.pth")
+        model.load_checkpoint(mf_param)
         model.Qformer.bert = torch.load('model/mf2_bert.pth', map_location=torch.device('cpu'))
         model.to('cuda')
     elif type == 'Biological Process':
-        model.load_checkpoint("model/checkpoint_bp1.pth")
+        # model.load_checkpoint("model/checkpoint_bp1.pth")
+        model.load_checkpoint(bp_param)
         model.Qformer.bert = torch.load('model/bp1_bert.pth', map_location=torch.device('cpu'))
         model.to('cuda')
     elif type == 'Cellar Component':
-        model.load_checkpoint("model/checkpoint_cc2.pth")
+        # model.load_checkpoint("model/checkpoint_cc2.pth")
+        model.load_checkpoint(cc_param)
         model.Qformer.bert = torch.load('model/cc2_bert.pth', map_location=torch.device('cpu'))
         model.to('cuda')
     return model
 
-
 models = {
     'Molecule Function': get_model('Molecule Function'),
     'Biological Process': get_model('Biological Process'),