Update app.py

app.py

@@ -32,12 +32,12 @@ def get_model(type='Molecule Function'):
 models = {
     'Molecule Function': get_model('Molecule Function'),
     'Biological Process': get_model('Biological Process'),
-    'Cellar Component': get_model('Cellar Component'),
+    'Cellular Component': get_model('Cellar Component'),
     }
 
 
 # Load the mistral model
-mistral_model = MistralForCausalLM.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B", torch_dtype=torch.float16)
+mistral_model = MistralForCausalLM.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B", torch_dtype=torch.float16).to('cuda')
 
 # Load ESM2 model
 model_esm, alphabet = pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D')
@@ -54,14 +54,23 @@ go_des['text'] = go_des['text'].apply(lambda x: x.lower())
 GO_dict = dict(zip(go_des['text'], go_des['id']))
 Func_dict = dict(zip(go_des['id'], go_des['text']))
 
-# terms_mf = pd.read_pickle('/cluster/home/wenkai/deepgo2/data/mf/terms.pkl')
 terms_mf = pd.read_pickle('data/terms/mf_terms.pkl')
 choices_mf = [Func_dict[i] for i in list(set(terms_mf['gos']))]
-choices = {x.lower(): x for x in choices_mf}
-
+choices_mf = {x.lower(): x for x in choices_mf}
+terms_bp = pd.read_pickle('data/terms/bp_terms.pkl')
+choices_bp = [Func_dict[i] for i in list(set(terms_bp['gos']))]
+choices_bp = {x.lower(): x for x in choices_bp}
+terms_cc = pd.read_pickle('data/terms/cc_terms.pkl')
+choices_cc = [Func_dict[i] for i in list(set(terms_cc['gos']))]
+choices_cc = {x.lower(): x for x in choices_cc}
+choices = {
+    'Molecule Function': choices_mf,
+    'Biological Process': choices_bp,
+    'Cellular Component': choices_cc,
+    }
 
 @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,36 +153,40 @@ def generate_caption(model_id, protein, prompt):
                'text_input': ['none'],
                'prompt': [prompt]}
 
-    model = models[model_id]
-    # Generate the output
-    prediction = model.generate(mistral_model, samples, length_penalty=0., num_beams=15, num_captions=10, temperature=1.,
-                                repetition_penalty=1.0)
-
-    x = prediction[0]
-    x = [eval(i) for i in x.split('; ')]
-    pred_terms = []
-    temp = []
-    for i in x:
-        txt = i[0]
-        prob = i[1]
-        sim_list = difflib.get_close_matches(txt.lower(), choices, n=1, cutoff=0.9)
-        if len(sim_list) > 0:
-            t_standard = sim_list[0]
-            if t_standard not in temp:
-                pred_terms.append(t_standard+f'({prob})')
-                temp.append(t_standard)
+    union_pred_terms = []
+    for model_id in models.keys():
+        model = models[model_id]
+        # Generate the output
+        prediction = model.generate(mistral_model, samples, length_penalty=0., num_beams=15, num_captions=10, temperature=1.,
+                                    repetition_penalty=1.0)
+        x = prediction[0]
+        x = [eval(i) for i in x.split('; ')]
+        pred_terms = []
+        temp = []
+        for i in x:
+            txt = i[0]
+            prob = i[1]
+            sim_list = difflib.get_close_matches(txt.lower(), choices[model_id], n=1, cutoff=0.9)
+            if len(sim_list) > 0:
+                t_standard = sim_list[0]
+                if t_standard not in temp:
+                    pred_terms.append(t_standard+f'({prob})')
+                    temp.append(t_standard)
+        union_pred_terms.append(pred_terms)
+
     if prompt == 'none':
         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 use other two types of model or try another sequence!"
-    if len(pred_terms) == 0:
+    if len(union_pred_terms[0]) == 0 and len(union_pred_terms[1]) == 0 and len(union_pred_terms[2]) == 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 = ''
+    if len(union_pred_terms[0]) != 0:
+        res_str += f"Based on the given amino acid sequence, the protein appears to have a primary function of {', '.join(pred_terms)}. "
+    if len(union_pred_terms[1]) != 0:
+        res_str += f"It is likely involved in the {', '.join(pred_terms)}. "
+    if len(union_pred_terms[2]) != 0:
+        res_str += f"It's subcellular localization is within the {', '.join(pred_terms)}."
     return res_str
     # return "test"
 
@@ -205,7 +218,6 @@ 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")
                 prompt = gr.Textbox(type="text", label="Taxonomy Prompt (Optional)")
                 submit_btn = gr.Button(value="Submit")
@@ -214,20 +226,20 @@ with gr.Blocks(css=css) as demo:
         # O14813 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)