Update app.py

app.py

@@ -6,12 +6,14 @@ import numpy as np
 from collections import Counter
 
 # --- Download model from HF Hub ---
-repo_id = "Ym420/Peptide-Function"  # replace with your HF repo
+repo_id = "Ym420/Peptide-Function"
 model_filename = "xgb_multilabel_model_full.pkl"
 
 model_path = hf_hub_download(repo_id=repo_id, filename=model_filename)
 model_package = joblib.load(model_path)
-model = model_package['model']
+
+# --- Unwrap model dict ---
+model_dict = model_package['model']           # dict: {'Gram+': XGBClassifier, ...}
 feature_columns = model_package['feature_columns']
 
 # --- Metadata ---
@@ -34,12 +36,12 @@ TARGET_CELLS = ["Gram+", "Fungus", "Mammalian Cell", "Cancer", "Gram-"]
 def extract_features_app(seq: str) -> pd.DataFrame:
     seq = seq.upper()
     
-    # --- 1. Dipeptide composition ---
+    # Dipeptide composition
     count = Counter([seq[i:i+2] for i in range(len(seq)-1)])
     total = max(len(seq)-1, 1)
     dipep_features = [count.get(dp, 0) / total for dp in dipeptides]
 
-    # --- 2. Physicochemical features ---
+    # Physicochemical features
     def g(aa, table): return table.get(aa, 0)
     def h(dp, table): return (g(dp[0], table) + g(dp[1], table)) / 2.0
 
@@ -67,7 +69,7 @@ def extract_features_app(seq: str) -> pd.DataFrame:
     features = dipep_features + physchem_features
     
     df = pd.DataFrame([features], columns=feature_columns)
-    df = df.astype('float32')  # ensure same type as training
+    df = df.astype('float32')
     return df
 
 # --- Prediction function ---
@@ -78,16 +80,17 @@ def predict_peptide(sequence: str):
 
     X = extract_features_app(seq)
     
-    probs_list = model.predict_proba(X)  # list of arrays per target
-    # --- probs_list: list of arrays from each estimator ---
-    #probs_list = [est.predict_proba(X) for est in model]  # model is a list
-
     table = []
-    for i, target in enumerate(TARGET_CELLS):
-        prob = float(probs_list[i][0][1])
-        table.append([target, round(prob, 4)])
-    return table
+    # Iterate over each target classifier
+    for target in TARGET_CELLS:
+        clf = model_dict.get(target)
+        if clf is not None:
+            prob = clf.predict_proba(X)[0][1]  # positive-class probability (0-1)
+            table.append([target, round(float(prob), 4)])
+        else:
+            table.append([target, None])
 
+    return table
 
 # --- Gradio Interface ---
 custom_css = """
@@ -115,8 +118,5 @@ with gr.Blocks(css=custom_css, theme="default") as demo:
     # API endpoint for iOS app
     gr.api(predict_peptide, api_name="predict_peptide")
 
-#if __name__ == "__main__": 
-#    demo.launch()
-
 if __name__ == "__main__": 
     demo.launch(show_error=True)