Update app.py

app.py

@@ -19,6 +19,7 @@ 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.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")
@@ -26,6 +27,7 @@ def get_model(type='Molecule Function'):
         model.to('cuda')
     elif type == 'Cellar Component':
         model.load_checkpoint("model/checkpoint_cc2.pth")
+        model.Qformer.bert = torch.load('model/cc2_bert.pth', map_location=torch.device('cpu'))
         model.to('cuda')
     return model