Update app.py

app.py

@@ -86,16 +86,16 @@ def generate_caption(protein, prompt):
     repr_layers = [36]
     truncation_seq_length = 1024
     toks_per_batch = 4096
-    print("start")
+    # print("start")
     dataset = FastaBatchedDataset([protein_name], [protein_seq])
-    print("dataset prepared")
+    # print("dataset prepared")
     batches = dataset.get_batch_indices(toks_per_batch, extra_toks_per_seq=1)
-    print("batches prepared")
+    # print("batches prepared")
 
     data_loader = torch.utils.data.DataLoader(
         dataset, collate_fn=alphabet.get_batch_converter(truncation_seq_length), batch_sampler=batches
     )
-    print(f"Read sequences")
+    # print(f"Read sequences")
     return_contacts = "contacts" in include
 
     assert all(-(model_esm.num_layers + 1) <= i <= model_esm.num_layers for i in repr_layers)
@@ -142,7 +142,7 @@ def generate_caption(protein, prompt):
         outputs = model_esm(**inputs)
     esm_emb = outputs.last_hidden_state.detach()[0]
     '''
-    print("esm embedding generated")
+    # print("esm embedding generated")
     esm_emb = F.pad(esm_emb.t(), (0, 1024 - len(esm_emb))).t().to('cuda')
     if prompt is None:
         prompt = 'none'
@@ -182,13 +182,15 @@ def generate_caption(protein, prompt):
         return res_str
     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(union_pred_terms[0])}. "
+        temp = ['- '+i+'\n' for i in union_pred_terms[0]]
+        res_str += f"Based on the given amino acid sequence, the protein appears to have a primary function of \n{temp} \n"
     if len(union_pred_terms[1]) != 0:
-        res_str += f"It is likely involved in the {', '.join(union_pred_terms[1])}. "
+        temp = ['- ' + i + '\n' for i in union_pred_terms[1]]
+        res_str += f"It is likely involved in the following process: \n{temp} \n"
     if len(union_pred_terms[2]) != 0:
-        res_str += f"It's subcellular localization is within the {', '.join(union_pred_terms[2])}."
+        temp = ['- ' + i + '\n' for i in union_pred_terms[2]]
+        res_str += f"It's subcellular localization is within the: \n{temp}"
     return res_str
-    # return "test"
 
 
 # Define the FAPM interface
@@ -226,7 +228,9 @@ with gr.Blocks(css=css) as demo:
                 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")
+                # output_text = gr.Textbox(label="Output Text")
+                with gr.Accordion('Answer:', open=True):
+                    output_markdown = gr.Markdown(label="Output")
         # O14813 train index 127, 266, 738, 1060 test index 4
         gr.Examples(
             examples=[
@@ -238,12 +242,12 @@ with gr.Blocks(css=css) as demo:
                 ['MNDLMIQLLDQFEMGLRERAIKVMATINDEKHRFPMELNKKQCSLMLLGTTDTTTFDMRFNSKKDFPRIKGAREKYPRDAVIEWYHQNWMRTEVKQ', 'Bacteriophage'],
             ],
             inputs=[input_protein, prompt],
-            outputs=[output_text],
+            outputs=[output_markdown],
             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], [output_markdown])
 
 demo.launch(debug=True)