feat: choose model

app.py

@@ -130,29 +130,49 @@ def process_fasta(fasta_text):
 
 
 # Create a Gradio interface
-def gradio_process_fasta(fasta_text):
+def predict(choice, fasta_text):
     """
     Wrapper for Gradio to process the FASTA text.
     """
     fasta_dict = process_fasta(fasta_text)
-    result = predictor.predict(fasta_dict)
-    new_result = {}
-    for key, value in result.items():
-        if value[1] == 1:
-            new_result[key] = {"probability": value[0], "class": "bitter"}
-        else:
-            new_result[key] = {"probability": 1 - value[0], "class": "non-bitter"}
-    return new_result
+    result = predictor.predict(fasta_dict, model_type=choice)
+    result_df = pd.DataFrame(
+        {
+            "id": list(result.keys()),
+            "probability": [i[0] for i in result.values()],
+            "class": ["bitter" if i[1] == 1 else "non-bitter" for i in result.values()],
+        }
+    )
+    # text_result = f"ID\tClass\tProbability\n"
+    # for key, value in result.items():
+    #     text_result += (
+    #         f"{key}\t{'bitter' if value[1] == 1 else 'non-bitter'}\t{value[0]}\n"
+    #     )
+    return result_df
 
 
 interface = gr.Interface(
-    fn=gradio_process_fasta,
-    inputs=gr.Textbox(
-        label="Enter FASTA format text", lines=10, placeholder=">id1\nPA\n>KA\nPPP"
+    fn=predict,
+    inputs=[
+        gr.Dropdown(
+            choices=["xBitterT5-640", "xBitterT5-720"],
+            label="Select xBitterT5 variant",
+            value="xBitterT5-720",
+        ),
+        gr.Textbox(
+            label="Enter peptide sequences in FASTA format",
+            lines=10,
+            placeholder=">id1\nVAPFPE\n>id2\nRRPP\n>id3\nGH\nid4\nGVDTK",
+        ),
+    ],
+    # outputs=gr.Textbox(label="Predictions", type="text"),
+    outputs=gr.Dataframe(
+        headers=["ID", "Class", "Probability"],
+        show_fullscreen_button=True,
     ),
-    outputs=gr.JSON(label="Prediction results"),
     title="xBitterT5",
-    description=("Enter a FASTA-formatted text"),
+    description=("Prediction of bitter peptides using xBitterT5."),
+    flagging_mode="never",
 )
 # Launch the Gradio app
 interface.launch()