Update app.py

app.py

@@ -7,17 +7,39 @@ from lavis.models.protein_models.protein_function_opt import Blip2ProteinMistral
 from lavis.models.base_model import FAPMConfig
 import spaces
 import gradio as gr
+# from esm_scripts.extract import run_demo
 from esm import pretrained, FastaBatchedDataset
 from data.evaluate_data.utils import Ontology
 import difflib
 import re
-
-
-# Load the model
-model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
-model.load_checkpoint("model/checkpoint_mf2.pth")
-model.to('cuda')
-
+from transformers import MistralForCausalLM
+
+# 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.to('cuda')
+    elif type == 'Biological Process':
+        model.load_checkpoint("model/checkpoint_bp1.pth")
+        model.to('cuda')
+    elif type == 'Cellar Component':
+        model.load_checkpoint("model/checkpoint_cc2.pth")
+        model.to('cuda')
+    return model
+
+
+models = {
+    'Molecule Function': get_model('Molecule Function'),
+    'Biological Process': get_model('Biological Process'),
+    'Cellar Component': get_model('Cellar Component'),
+    }
+
+
+# Load the mistral model
+mistral_model = MistralForCausalLM.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B", torch_dtype=torch.float16)
+
+# Load ESM2 model
 model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
 model_esm.to('cuda')
 model_esm.eval()
@@ -39,7 +61,7 @@ choices = {x.lower(): x for x in choices_mf}
 
 
 @spaces.GPU
-def generate_caption(protein, prompt):
+def generate_caption(model_id, protein, prompt):
     # Process the image and the prompt
     # with open('/home/user/app/example.fasta', 'w') as f:
     #     f.write('>{}\n'.format("protein_name"))
@@ -122,8 +144,9 @@ def generate_caption(protein, prompt):
                'text_input': ['none'],
                'prompt': [prompt]}
 
+    model = models[model_id]
     # Generate the output
-    prediction = model.generate(samples, length_penalty=0., num_beams=15, num_captions=10, temperature=1.,
+    prediction = model.generate(mistral_model, samples, length_penalty=0., num_beams=15, num_captions=10, temperature=1.,
                                 repetition_penalty=1.0)
 
     x = prediction[0]
@@ -140,12 +163,17 @@ def generate_caption(protein, prompt):
                 pred_terms.append(t_standard+f'({prob})')
                 temp.append(t_standard)
     if prompt == 'none':
-        res_str = "No available predictions for this protein, you can try to remove prompt!"
+        res_str = "No available predictions for this protein, you can use other two types of model, remove prompt or try another sequence!"
     else:
-        res_str = "No available predictions for this protein, you can try another protein sequence!"
+        res_str = "No available predictions for this protein, you can use other two types of model or try another sequence!"
     if len(pred_terms) == 0:
         return res_str
-    res_str = f"Based on the given amino acid sequence, the protein appears to have a primary function of {', '.join(pred_terms)}"
+    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"
 
@@ -155,7 +183,6 @@ 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")],
@@ -165,7 +192,6 @@ The model used in this app is available at [Hugging Face Model Hub](https://hugg
 # # Launch the interface
 # iface.launch()
 
-
 css = """
   #output {
     height: 500px; 
@@ -179,30 +205,29 @@ with gr.Blocks(css=css) as demo:
     with gr.Tab(label="Protein caption"):
         with gr.Row():
             with gr.Column():
+                model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value='Molecule Function')
                 input_protein = gr.Textbox(type="text", label="Upload sequence")
-                # model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value='microsoft/Florence-2-large')
                 prompt = gr.Textbox(type="text", label="Taxonomy Prompt (Optional)")
                 submit_btn = gr.Button(value="Submit")
             with gr.Column():
                 output_text = gr.Textbox(label="Output Text")
-        # train index 127, 266, 738, 1060 test index 4
+        # O14813 train index 127, 266, 738, 1060 test index 4
         gr.Examples(
             examples=[
-                ["MDYSYLNSYDSCVAAMEASAYGDFGACSQPGGFQYSPLRPAFPAAGPPCPALGSSNCALGALRDHQPAPYSAVPYKFFPEPSGLHEKRKQRRIRTTFTSAQLKELERVFAETHYPDIYTREELALKIDLTEARVQVWFQNRRAKFRKQERAASAKGAAGAAGAKKGEARCSSEDDDSKESTCSPTPDSTASLPPPPAPGLASPRLSPSPLPVALGSGPGPGPGPQPLKGALWAGVAGGGGGGPGAGAAELLKAWQPAESGPGPFSGVLSSFHRKPGPALKTNLF", ''],
-                ["MKTLALFLVLVCVLGLVQSWEWPWNRKPTKFPIPSPNPRDKWCRLNLGPAWGGRC", ''],
-                ["MAAAGGARLLRAASAVLGGPAGRWLHHAGSRAGSSGLLRNRGPGGSAEASRSLSVSARARSSSEDKITVHFINRDGETLTTKGKVGDSLLDVVVENNLDIDGFGACEGTLACSTCHLIFEDHIYEKLDAITDEENDMLDLAYGLTDRSRLGCQICLTKSMDNMTVRVPETVADARQSIDVGKTS", 'Homo'],
-                ['MASAELSREENVYMAKLAEQAERYEEMVEFMEKVAKTVDSEELTVEERNLLSVAYKNVIGARRASWRIISSIEQKEEGRGNEDRVTLIKDYRGKIETELTKICDGILKLLETHLVPSSTAPESKVFYLKMKGDYYRYLAEFKTGAERKDAAENTMVAYKAAQDIALAELAPTHPIRLGLALNFSVFYYEILNSPDRACSLAKQAFDEAISELDTLSEESYKDSTLIMQLLRDNLTLWTSDISEDPAEEIREAPKRDSSEGQ', 'Zea'],
-                ['MIKAAVTKESLYRMNTLMEAFQGFLGLDLGEFTFKVKPGVFLLTDVKSYLIGDKYDDAFNALIDFVLRNDRDAVEGTETDVSIRLGLSPSDMVVKRQDKTFTFTHGDLEFEVHWINL', 'Bacteriophage'],
-                ['MNDLMIQLLDQFEMGLRERAIKVMATINDEKHRFPMELNKKQCSLMLLGTTDTTTFDMRFNSKKDFPRIKGAREKYPRDAVIEWYHQNWMRTEVKQ', 'Bacteriophage'],
+                ["Molecule Function", "MDYSYLNSYDSCVAAMEASAYGDFGACSQPGGFQYSPLRPAFPAAGPPCPALGSSNCALGALRDHQPAPYSAVPYKFFPEPSGLHEKRKQRRIRTTFTSAQLKELERVFAETHYPDIYTREELALKIDLTEARVQVWFQNRRAKFRKQERAASAKGAAGAAGAKKGEARCSSEDDDSKESTCSPTPDSTASLPPPPAPGLASPRLSPSPLPVALGSGPGPGPGPQPLKGALWAGVAGGGGGGPGAGAAELLKAWQPAESGPGPFSGVLSSFHRKPGPALKTNLF", ''],
+                ["Molecule Function", "MKTLALFLVLVCVLGLVQSWEWPWNRKPTKFPIPSPNPRDKWCRLNLGPAWGGRC", ''],
+                ["Molecule Function", "MAAAGGARLLRAASAVLGGPAGRWLHHAGSRAGSSGLLRNRGPGGSAEASRSLSVSARARSSSEDKITVHFINRDGETLTTKGKVGDSLLDVVVENNLDIDGFGACEGTLACSTCHLIFEDHIYEKLDAITDEENDMLDLAYGLTDRSRLGCQICLTKSMDNMTVRVPETVADARQSIDVGKTS", 'Homo'],
+                ["Molecule Function", 'MASAELSREENVYMAKLAEQAERYEEMVEFMEKVAKTVDSEELTVEERNLLSVAYKNVIGARRASWRIISSIEQKEEGRGNEDRVTLIKDYRGKIETELTKICDGILKLLETHLVPSSTAPESKVFYLKMKGDYYRYLAEFKTGAERKDAAENTMVAYKAAQDIALAELAPTHPIRLGLALNFSVFYYEILNSPDRACSLAKQAFDEAISELDTLSEESYKDSTLIMQLLRDNLTLWTSDISEDPAEEIREAPKRDSSEGQ', 'Zea'],
+                ["Molecule Function", 'MIKAAVTKESLYRMNTLMEAFQGFLGLDLGEFTFKVKPGVFLLTDVKSYLIGDKYDDAFNALIDFVLRNDRDAVEGTETDVSIRLGLSPSDMVVKRQDKTFTFTHGDLEFEVHWINL', 'Bacteriophage'],
+                ["Molecule Function", 'MNDLMIQLLDQFEMGLRERAIKVMATINDEKHRFPMELNKKQCSLMLLGTTDTTTFDMRFNSKKDFPRIKGAREKYPRDAVIEWYHQNWMRTEVKQ', 'Bacteriophage'],
             ],
-            inputs=[input_protein, prompt],
+            inputs=[model_selector, input_protein, prompt],
             outputs=[output_text],
             fn=generate_caption,
             cache_examples=True,
             label='Try examples'
         )
-
-        submit_btn.click(generate_caption, [input_protein, prompt], [output_text])
+        submit_btn.click(generate_caption, [model_selector, input_protein, prompt], [output_text])
 
 demo.launch(debug=True)