Update app.py

app.py

@@ -30,6 +30,7 @@ def get_model(type='Molecule Function'):
     elif type == 'Cellar Component':
         model.load_checkpoint("model/checkpoint_cc2.pth")
         model.to('cuda')
+    return model
 
 
 models = {
@@ -161,8 +162,15 @@ def generate_caption(model_id, protein, prompt):
             if t_standard not in temp:
                 pred_terms.append(t_standard+f'({prob})')
                 temp.append(t_standard)
-
-    res_str = f"Based on the given amino acid sequence, the proteinappears to have a primary function of {', '.join(pred_terms)}"
+    res_str = "No available predictions for this protein! You can try the other two types of model or remove prompt."
+    if len(pred_terms) == 0:
+        return res_str
+    if model_id == 'Molecule Function':
+        res_str = f"Based on the given amino acid sequence, the protein appears to have a primary function of {', '.join(pred_terms)}"
+    elif model_id == 'Biological Process':
+        res_str = f"Based on the given amino acid sequence, it is likely involved in the {', '.join(pred_terms)}"
+    elif model_id == 'Cellar Component':
+        res_str = f"Based on the given amino acid sequence, it's subcellular localization is within the {', '.join(pred_terms)}"
     return res_str
     # return "test"