Update app.py

app.py

@@ -51,6 +51,9 @@ def extract_features(sequence: str) -> pd.DataFrame:
     return df
 
 # --- Prediction function ---
+TARGET_CELLS = ["Gram+", "Fungus", "Mammalian Cell", "Cancer", "Gram-"]
+
+
 def predict_peptide(sequence: str):
     seq = "".join(sequence.split()).upper()
     if not seq:
@@ -59,12 +62,14 @@ def predict_peptide(sequence: str):
     X = extract_features(seq)
     probs_list = model.predict_proba(X)  # list of arrays per target cell
 
-    # Format output as table: Target Cell | Probability
+    # Format output with 4 decimal places
     table = []
-    for i, target in enumerate(model.classes_):
-        table.append([target, float(probs_list[i][0][1])])
+    for i, target in enumerate(TARGET_CELLS):
+        prob = float(probs_list[i][0][1])
+        table.append([target, round(prob, 4)])
     return table
 
+
 # --- Gradio Interface ---
 custom_css = """
 footer, .footer {display:none !important;}