Update app.py

app.py

@@ -4,12 +4,27 @@ import numpy as np
 import pandas as pd
 from propy import AAComposition, Autocorrelation, CTD, PseudoAAC
 from sklearn.preprocessing import MinMaxScaler
+import torch
+from transformers import BertTokenizer, BertModel
+from math import expm1
 
-# Load model and scaler
+# =====================
+# Load AMP Classifier
+# =====================
 model = joblib.load("RF.joblib")
 scaler = joblib.load("norm (4).joblib")
 
-# Feature list (KEEP THIS CONSISTENT)
+# =====================
+# Load ProtBert Globally
+# =====================
+tokenizer = BertTokenizer.from_pretrained("Rostlab/prot_bert", do_lower_case=False)
+protbert_model = BertModel.from_pretrained("Rostlab/prot_bert")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+protbert_model = protbert_model.to(device).eval()
+
+# =====================
+# Feature List (ProPy)
+# =====================
 selected_features = [
     "_SolventAccessibilityC3", "_SecondaryStrC1", "_SecondaryStrC3", "_ChargeC1", "_PolarityC1",
     "_NormalizedVDWVC1", "_HydrophobicityC3", "_SecondaryStrT23", "_PolarizabilityD1001",
@@ -45,16 +60,17 @@ selected_features = [
     "APAAC24"
 ]
 
+# =====================
+# AMP Feature Extractor
+# =====================
 def extract_features(sequence):
-
     all_features_dict = {}
+    sequence = ''.join([aa for aa in sequence.upper() if aa in "ACDEFGHIKLMNPQRSTVWY"])
+    if len(sequence) < 10:
+        return "Error: Sequence too short."
 
-      # Calculate all dipeptide features
     dipeptide_features = AAComposition.CalculateAADipeptideComposition(sequence)
-    
-    # Add only the first 420 features to the dictionary
-    first_420_keys = list(dipeptide_features.keys())[:420]  # Get the first 420 keys
-    filtered_dipeptide_features = {key: dipeptide_features[key] for key in first_420_keys}
+    filtered_dipeptide_features = {k: dipeptide_features[k] for k in list(dipeptide_features.keys())[:420]}
     ctd_features = CTD.CalculateCTD(sequence)
     auto_features = Autocorrelation.CalculateAutoTotal(sequence)
     pseudo_features = PseudoAAC.GetAPseudoAAC(sequence, lamda=9)
@@ -64,65 +80,45 @@ def extract_features(sequence):
     all_features_dict.update(auto_features)
     all_features_dict.update(pseudo_features)
 
-    # Convert all features to DataFrame
     feature_df_all = pd.DataFrame([all_features_dict])
+    normalized_array = scaler.transform(feature_df_all.values)
+    normalized_df = pd.DataFrame(normalized_array, columns=feature_df_all.columns)
+    selected_df = normalized_df[selected_features].fillna(0)
 
-    # Normalize ALL features
-    normalized_feature_array = scaler.transform(feature_df_all.values) # Normalize the numpy array
-    normalized_feature_df = pd.DataFrame(normalized_feature_array, columns=feature_df_all.columns) # Convert back to DataFrame with original column names
-
-    # Select features AFTER normalization
-    feature_df_selected = normalized_feature_df[selected_features].copy()
-    feature_df_selected = feature_df_selected.fillna(0) # Fill missing if any after selection (though unlikely now)
-    feature_array = feature_df_selected.values
-
-
-    return feature_array
-
+    return selected_df.values
 
+# =====================
+# AMP Classifier
+# =====================
 def predict(sequence):
-    """Predicts whether the input sequence is an AMP."""
     features = extract_features(sequence)
-    if isinstance(features, str) and features.startswith("Error:"):
+    if isinstance(features, str):
         return features
-
     prediction = model.predict(features)[0]
     probabilities = model.predict_proba(features)[0]
-
     if prediction == 0:
         return f"{probabilities[0] * 100:.2f}% chance of being an Antimicrobial Peptide (AMP)"
     else:
         return f"{probabilities[1] * 100:.2f}% chance of being Non-AMP"
 
-        
+# =====================
+# MIC Predictor (ProtBert-based)
+# =====================
 def predictmic(sequence):
-    import torch
-    from transformers import BertTokenizer, BertModel
-    import numpy as np
-    import joblib  # ✅ Use joblib instead of pickle
-    from math import expm1
-
-    # === Load ProtBert model ===
-    tokenizer = BertTokenizer.from_pretrained("Rostlab/prot_bert", do_lower_case=False)
-    model = BertModel.from_pretrained("Rostlab/prot_bert")
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = model.to(device).eval()
-
-    # === Preprocess input sequence ===
     sequence = ''.join([aa for aa in sequence.upper() if aa in "ACDEFGHIKLMNPQRSTVWY"])
     if len(sequence) < 10:
         return {"Error": "Sequence too short or invalid. Must contain at least 10 valid amino acids."}
 
-    # === Tokenize & embed using mean pooling ===
+    # Tokenize
     seq_spaced = ' '.join(list(sequence))
     tokens = tokenizer(seq_spaced, return_tensors="pt", padding='max_length', truncation=True, max_length=512)
     tokens = {k: v.to(device) for k, v in tokens.items()}
 
     with torch.no_grad():
-        outputs = model(**tokens)
-        embedding = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy().reshape(1, -1)  # Shape: (1, 1024)
+        outputs = protbert_model(**tokens)
+        embedding = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy().reshape(1, -1)
 
-    # === MIC models and scalers for each bacterium ===
+    # MIC model config
     bacteria_config = {
         "E.coli": {
             "model": "coli_xgboost_model.pkl",
@@ -142,59 +138,46 @@ def predictmic(sequence):
         "K.Pneumonia": {
             "model": "pne_mlp_model.pkl",
             "scaler": "pne_scaler.pkl",
-            "pca": "pne_pca"  # Ensure this PCA file was also saved using joblib
+            "pca": "pne_pca"
         }
     }
 
     mic_results = {}
-
     for bacterium, cfg in bacteria_config.items():
         try:
-            # === Load scaler and transform ===
             scaler = joblib.load(cfg["scaler"])
             scaled = scaler.transform(embedding)
-
-            # === Apply PCA if exists ===
-            if cfg["pca"] is not None:
+            if cfg["pca"]:
                 pca = joblib.load(cfg["pca"])
                 transformed = pca.transform(scaled)
             else:
                 transformed = scaled
-
-            # === Load model and predict ===
-            mic_model = joblib.load(cfg["model"])
-            mic_log = mic_model.predict(transformed)[0]
-            mic = round(expm1(mic_log), 3)  # Inverse of log1p used in training
-
+            model = joblib.load(cfg["model"])
+            mic_log = model.predict(transformed)[0]
+            mic = round(expm1(mic_log), 3)
             mic_results[bacterium] = mic
-
         except Exception as e:
             mic_results[bacterium] = f"Error: {str(e)}"
 
     return mic_results
 
-
-
-
-
+# =====================
+# Combined Prediction Function
+# =====================
 def full_prediction(sequence):
-    # AMP prediction
     features = extract_features(sequence)
-    if isinstance(features, str) and features.startswith("Error:"):
-        return "Error", 0.0, {}
-
+    if isinstance(features, str):
+        return "Error", "0%", {}
     prediction = model.predict(features)[0]
     probabilities = model.predict_proba(features)[0]
-
     amp_result = "Antimicrobial Peptide (AMP)" if prediction == 0 else "Non-AMP"
     confidence = round(probabilities[0 if prediction == 0 else 1] * 100, 2)
-
-    # MIC prediction
     mic_values = predictmic(sequence)
-
     return amp_result, f"{confidence}%", mic_values
 
-import gradio as gr
+# =====================
+# Gradio Interface
+# =====================
 iface = gr.Interface(
     fn=full_prediction,
     inputs=gr.Textbox(label="Enter Protein Sequence"),
@@ -204,7 +187,7 @@ iface = gr.Interface(
         gr.JSON(label="Predicted MIC (µg/mL) for Each Bacterium")
     ],
     title="AMP & MIC Predictor",
-    description="Enter an amino acid sequence (e.g., FLPVLAGGL) to predict AMP class and MIC values across bacteria."
+    description="Enter an amino acid sequence (≥10 valid letters) to predict AMP class and MIC values."
 )
 
 iface.launch(share=True)