Update app.py

app.py

@@ -7,7 +7,6 @@ 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
@@ -15,30 +14,9 @@ import re
 
 
 # Load the model
-# model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
-# model.load_checkpoint("model/checkpoint_mf2.pth")
-# model.to('cuda')
-
-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'),
-    }
-
+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')
@@ -61,7 +39,7 @@ choices = {x.lower(): x for x in choices_mf}
 
 
 @spaces.GPU
-def generate_caption(model_id, protein, prompt):
+def generate_caption(protein, prompt):
     # Process the image and the prompt
     # with open('/home/user/app/example.fasta', 'w') as f:
     #     f.write('>{}\n'.format("protein_name"))
@@ -144,7 +122,6 @@ def generate_caption(model_id, 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.,
                                 repetition_penalty=1.0)
@@ -162,15 +139,13 @@ 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 = "No available predictions for this protein! You can try the other two types of model or remove prompt."
+    if prompt == 'none':
+        res_str = "No available predictions for this protein, you can try to remove prompt!"
+    else:
+        res_str = "No available predictions for this protein, you can try another protein sequence!"
     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)}"
+    res_str = f"Based on the given amino acid sequence, the protein appears to have a primary function of {', '.join(pred_terms)}"
     return res_str
     # return "test"
 
@@ -180,6 +155,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")],
@@ -189,6 +165,7 @@ The model used in this app is available at [Hugging Face Model Hub](https://hugg
 # # Launch the interface
 # iface.launch()
 
+
 css = """
   #output {
     height: 500px; 
@@ -202,29 +179,30 @@ 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")
-        # O14813 train index 127, 266, 738, 1060 test index 4
+        # train index 127, 266, 738, 1060 test index 4
         gr.Examples(
             examples=[
-                ["Molecule Function", "MDYSYLNSYDSCVAAMEASAYGDFGACSQPGGFQYSPLRPAFPAAGPPCPALGSSNCALGALRDHQPAPYSAVPYKFFPEPSGLHEKRKQRRIRTTFTSAQLKELERVFAETHYPDIYTREELALKIDLTEARVQVWFQNRRAKFRKQERAASAKGAAGAAGAKKGEARCSSEDDDSKESTCSPTPDSTASLPPPPAPGLASPRLSPSPLPVALGSGPGPGPGPQPLKGALWAGVAGGGGGGPGAGAAELLKAWQPAESGPGPFSGVLSSFHRKPGPALKTNLF", ''],
-                ["Molecule Function", "MKTLALFLVLVCVLGLVQSWEWPWNRKPTKFPIPSPNPRDKWCRLNLGPAWGGRC", ''],
-                ["Molecule Function", "MAAAGGARLLRAASAVLGGPAGRWLHHAGSRAGSSGLLRNRGPGGSAEASRSLSVSARARSSSEDKITVHFINRDGETLTTKGKVGDSLLDVVVENNLDIDGFGACEGTLACSTCHLIFEDHIYEKLDAITDEENDMLDLAYGLTDRSRLGCQICLTKSMDNMTVRVPETVADARQSIDVGKTS", 'Homo'],
-                ["Molecule Function", 'MASAELSREENVYMAKLAEQAERYEEMVEFMEKVAKTVDSEELTVEERNLLSVAYKNVIGARRASWRIISSIEQKEEGRGNEDRVTLIKDYRGKIETELTKICDGILKLLETHLVPSSTAPESKVFYLKMKGDYYRYLAEFKTGAERKDAAENTMVAYKAAQDIALAELAPTHPIRLGLALNFSVFYYEILNSPDRACSLAKQAFDEAISELDTLSEESYKDSTLIMQLLRDNLTLWTSDISEDPAEEIREAPKRDSSEGQ', 'Zea'],
-                ["Molecule Function", 'MIKAAVTKESLYRMNTLMEAFQGFLGLDLGEFTFKVKPGVFLLTDVKSYLIGDKYDDAFNALIDFVLRNDRDAVEGTETDVSIRLGLSPSDMVVKRQDKTFTFTHGDLEFEVHWINL', 'Bacteriophage'],
-                ["Molecule Function", 'MNDLMIQLLDQFEMGLRERAIKVMATINDEKHRFPMELNKKQCSLMLLGTTDTTTFDMRFNSKKDFPRIKGAREKYPRDAVIEWYHQNWMRTEVKQ', 'Bacteriophage'],
+                ["MDYSYLNSYDSCVAAMEASAYGDFGACSQPGGFQYSPLRPAFPAAGPPCPALGSSNCALGALRDHQPAPYSAVPYKFFPEPSGLHEKRKQRRIRTTFTSAQLKELERVFAETHYPDIYTREELALKIDLTEARVQVWFQNRRAKFRKQERAASAKGAAGAAGAKKGEARCSSEDDDSKESTCSPTPDSTASLPPPPAPGLASPRLSPSPLPVALGSGPGPGPGPQPLKGALWAGVAGGGGGGPGAGAAELLKAWQPAESGPGPFSGVLSSFHRKPGPALKTNLF", ''],
+                ["MKTLALFLVLVCVLGLVQSWEWPWNRKPTKFPIPSPNPRDKWCRLNLGPAWGGRC", ''],
+                ["MAAAGGARLLRAASAVLGGPAGRWLHHAGSRAGSSGLLRNRGPGGSAEASRSLSVSARARSSSEDKITVHFINRDGETLTTKGKVGDSLLDVVVENNLDIDGFGACEGTLACSTCHLIFEDHIYEKLDAITDEENDMLDLAYGLTDRSRLGCQICLTKSMDNMTVRVPETVADARQSIDVGKTS", 'Homo'],
+                ['MASAELSREENVYMAKLAEQAERYEEMVEFMEKVAKTVDSEELTVEERNLLSVAYKNVIGARRASWRIISSIEQKEEGRGNEDRVTLIKDYRGKIETELTKICDGILKLLETHLVPSSTAPESKVFYLKMKGDYYRYLAEFKTGAERKDAAENTMVAYKAAQDIALAELAPTHPIRLGLALNFSVFYYEILNSPDRACSLAKQAFDEAISELDTLSEESYKDSTLIMQLLRDNLTLWTSDISEDPAEEIREAPKRDSSEGQ', 'Zea'],
+                ['MIKAAVTKESLYRMNTLMEAFQGFLGLDLGEFTFKVKPGVFLLTDVKSYLIGDKYDDAFNALIDFVLRNDRDAVEGTETDVSIRLGLSPSDMVVKRQDKTFTFTHGDLEFEVHWINL', 'Bacteriophage'],
+                ['MNDLMIQLLDQFEMGLRERAIKVMATINDEKHRFPMELNKKQCSLMLLGTTDTTTFDMRFNSKKDFPRIKGAREKYPRDAVIEWYHQNWMRTEVKQ', 'Bacteriophage'],
             ],
-            inputs=[model_selector, input_protein, prompt],
+            inputs=[input_protein, prompt],
             outputs=[output_text],
             fn=generate_caption,
             cache_examples=True,
             label='Try examples'
         )
-        submit_btn.click(generate_caption, [model_selector, input_protein, prompt], [output_text])
+
+        submit_btn.click(generate_caption, [input_protein, prompt], [output_text])
 
 demo.launch(debug=True)