Update app.py

app.py

@@ -7,7 +7,7 @@ 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_scripts.extract import run_demo
 from esm import pretrained, FastaBatchedDataset
 from data.evaluate_data.utils import Ontology
 import difflib
@@ -15,9 +15,29 @@ import re
 
 
 # Load the model
-model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
-model.load_checkpoint("model/checkpoint_mf2.pth")
-model.to('cuda')
+# 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')
+
+
+models = {
+    'Molecule Function': get_model('Molecule Function'),
+    'Biological Process': get_model('Biological Process'),
+    'Cellar Component': get_model('Cellar Component'),
+    }
+
 
 model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
 model_esm.to('cuda')
@@ -40,7 +60,7 @@ choices = {x.lower(): x for x in choices_mf}
 
 
 @spaces.GPU
-def generate_caption(protein, prompt):
+def generate_caption(protein, prompt, model_id):
     # Process the image and the prompt
     # with open('/home/user/app/example.fasta', 'w') as f:
     #     f.write('>{}\n'.format("protein_name"))
@@ -123,6 +143,7 @@ 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.,
                                 repetition_penalty=1.0)
@@ -151,7 +172,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")],
@@ -161,7 +181,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; 
@@ -175,34 +194,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")
+                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 99, 127, 266, 738, 1060 test index 4
+        # O14813 train index 127, 266, 738, 1060 test index 4
         gr.Examples(
             examples=[
-                ["MKTLLLTLVVVTIVCLDLGNSLKCYVSREGKTQTCPEGEKLCEKYAVSYFHDGRWRYRYECTSACHRGPYNVCCSTDLCNK", 'Micrurus'],
-                ["MSSSAGSGHQPSQSRAIPTRTVAISDAAQLPHDYCTTPGGTLFSTTPGGTRIIYDRKFLLDRRNSPMAQTPPCHLPNIPGVTSPGTLIEDSKVEVNNLNNLNNHDRKHAVGDDAQFEMDI", 'Homo'],
-                ["MKTLALFLVLVCVLGLVQSWEWPWNRKPTKFPIPSPNPRDKWCRLNLGPAWGGRC", 'Sophophora'],
-                ["MAARGAMLRYLRVNVNPTIQNPRECVLPFSILLRRFSEEVRGSFLDKSEVTDRVLSVVKNFQKVDPSKVTPKANFQNDLGLDSLDSVEVVMALEEEFGFEIPDNEADKIQSIDLAVDFIASHPQAK", 'Arabidopsis'],
+                ["MDYSYLNSYDSCVAAMEASAYGDFGACSQPGGFQYSPLRPAFPAAGPPCPALGSSNCALGALRDHQPAPYSAVPYKFFPEPSGLHEKRKQRRIRTTFTSAQLKELERVFAETHYPDIYTREELALKIDLTEARVQVWFQNRRAKFRKQERAASAKGAAGAAGAKKGEARCSSEDDDSKESTCSPTPDSTASLPPPPAPGLASPRLSPSPLPVALGSGPGPGPGPQPLKGALWAGVAGGGGGGPGAGAAELLKAWQPAESGPGPFSGVLSSFHRKPGPALKTNLF", ''],
+                ["MKTLALFLVLVCVLGLVQSWEWPWNRKPTKFPIPSPNPRDKWCRLNLGPAWGGRC", ''],
                 ["MAAAGGARLLRAASAVLGGPAGRWLHHAGSRAGSSGLLRNRGPGGSAEASRSLSVSARARSSSEDKITVHFINRDGETLTTKGKVGDSLLDVVVENNLDIDGFGACEGTLACSTCHLIFEDHIYEKLDAITDEENDMLDLAYGLTDRSRLGCQICLTKSMDNMTVRVPETVADARQSIDVGKTS", 'Homo'],
                 ['MASAELSREENVYMAKLAEQAERYEEMVEFMEKVAKTVDSEELTVEERNLLSVAYKNVIGARRASWRIISSIEQKEEGRGNEDRVTLIKDYRGKIETELTKICDGILKLLETHLVPSSTAPESKVFYLKMKGDYYRYLAEFKTGAERKDAAENTMVAYKAAQDIALAELAPTHPIRLGLALNFSVFYYEILNSPDRACSLAKQAFDEAISELDTLSEESYKDSTLIMQLLRDNLTLWTSDISEDPAEEIREAPKRDSSEGQ', 'Zea'],
                 ['MIKAAVTKESLYRMNTLMEAFQGFLGLDLGEFTFKVKPGVFLLTDVKSYLIGDKYDDAFNALIDFVLRNDRDAVEGTETDVSIRLGLSPSDMVVKRQDKTFTFTHGDLEFEVHWINL', 'Bacteriophage'],
-                ['MQMYKLTAGTTGYHTLLTRTQAEHMLSLWGDKYSIDDCTPSNPIYSPSRYTKLELVYMAANATA', 'Bacteriophage'],
-                ['MSITAMDAKLQRILEESTCFGIGHDPNVKECKMCDVREQCKAKTQGMNVPTPTRKKPEDVAPAKEKPTTKKTTAKKSTAKEEKKETAPKAKETKAKPKSKPKKAKAPENPNLPNFKEMSFEELVELAKERNVEWKDYNSPNITRMRLIMALKASY', 'Bacteriophage'],
                 ['MNDLMIQLLDQFEMGLRERAIKVMATINDEKHRFPMELNKKQCSLMLLGTTDTTTFDMRFNSKKDFPRIKGAREKYPRDAVIEWYHQNWMRTEVKQ', 'Bacteriophage'],
             ],
-            inputs=[input_protein, prompt],
+            inputs=[input_protein, prompt, model_selector],
             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, [input_protein, prompt, model_selector], [output_text])
 
 demo.launch(debug=True)