Update app.py

app.py

@@ -19,35 +19,10 @@ 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()
 
-# Define selected features
-selected_features = ["_SolventAccessibilityC3", "_SecondaryStrC1", "_SecondaryStrC3", "_ChargeC1", "_PolarityC1",
-"_NormalizedVDWVC1", "_HydrophobicityC3", "_SecondaryStrT23", "_PolarizabilityD1001", "_PolarizabilityD2001",
-"_PolarizabilityD3001", "_SolventAccessibilityD1001", "_SolventAccessibilityD2001", "_SolventAccessibilityD3001",
-"_SecondaryStrD1001", "_SecondaryStrD1075", "_SecondaryStrD2001", "_SecondaryStrD3001", "_ChargeD1001",
-"_ChargeD1025", "_ChargeD2001", "_ChargeD3075", "_ChargeD3100", "_PolarityD1001", "_PolarityD1050",
-"_PolarityD2001", "_PolarityD3001", "_NormalizedVDWVD1001", "_NormalizedVDWVD2001", "_NormalizedVDWVD2025",
-"_NormalizedVDWVD2050", "_NormalizedVDWVD3001", "_HydrophobicityD1001", "_HydrophobicityD2001",
-"_HydrophobicityD3001", "_HydrophobicityD3025", "A", "R", "D", "C", "E", "Q", "H", "I", "M", "P", "Y", "V",
-"AR", "AV", "RC", "RL", "RV", "CR", "CC", "CL", "CK", "EE", "EI", "EL", "HC", "IA", "IL", "IV", "LA", "LC", "LE",
-"LI", "LT", "LV", "KC", "MA", "MS", "SC", "TC", "TV", "YC", "VC", "VE", "VL", "VK", "VV",
-"MoreauBrotoAuto_FreeEnergy30", "MoranAuto_Hydrophobicity2", "MoranAuto_Hydrophobicity4",
-"GearyAuto_Hydrophobicity20", "GearyAuto_Hydrophobicity24", "GearyAuto_Hydrophobicity26",
-"GearyAuto_Hydrophobicity27", "GearyAuto_Hydrophobicity28", "GearyAuto_Hydrophobicity29",
-"GearyAuto_Hydrophobicity30", "GearyAuto_AvFlexibility22", "GearyAuto_AvFlexibility26",
-"GearyAuto_AvFlexibility27", "GearyAuto_AvFlexibility28", "GearyAuto_AvFlexibility29", "GearyAuto_AvFlexibility30",
-"GearyAuto_Polarizability22", "GearyAuto_Polarizability24", "GearyAuto_Polarizability25",
-"GearyAuto_Polarizability27", "GearyAuto_Polarizability28", "GearyAuto_Polarizability29",
-"GearyAuto_Polarizability30", "GearyAuto_FreeEnergy24", "GearyAuto_FreeEnergy25", "GearyAuto_FreeEnergy30",
-"GearyAuto_ResidueASA21", "GearyAuto_ResidueASA22", "GearyAuto_ResidueASA23", "GearyAuto_ResidueASA24",
-"GearyAuto_ResidueASA30", "GearyAuto_ResidueVol21", "GearyAuto_ResidueVol24", "GearyAuto_ResidueVol25",
-"GearyAuto_ResidueVol26", "GearyAuto_ResidueVol28", "GearyAuto_ResidueVol29", "GearyAuto_ResidueVol30",
-"GearyAuto_Steric18", "GearyAuto_Steric21", "GearyAuto_Steric26", "GearyAuto_Steric27", "GearyAuto_Steric28",
-"GearyAuto_Steric29", "GearyAuto_Steric30", "GearyAuto_Mutability23", "GearyAuto_Mutability25",
-"GearyAuto_Mutability26", "GearyAuto_Mutability27", "GearyAuto_Mutability28", "GearyAuto_Mutability29",
-"GearyAuto_Mutability30", "APAAC1", "APAAC4", "APAAC5", "APAAC6", "APAAC8", "APAAC9", "APAAC12", "APAAC13",
-"APAAC15", "APAAC18", "APAAC19", "APAAC24"]
-
-# Create dummy data for LIME initialization
+# Define selected features (put your complete list here)
+selected_features = [ ... ]  # Replace with your full list
+
+# Dummy data for LIME
 sample_data = np.random.rand(100, len(selected_features))
 explainer = LimeTabularExplainer(
     training_data=sample_data,
@@ -56,7 +31,7 @@ explainer = LimeTabularExplainer(
     mode="classification"
 )
 
-# Feature extraction from sequence
+# Feature extraction function
 def extract_features(sequence):
     sequence = ''.join([aa for aa in sequence.upper() if aa in "ACDEFGHIKLMNPQRSTVWY"])
     if len(sequence) < 10:
@@ -87,7 +62,7 @@ def extract_features(sequence):
     except Exception as e:
         return f"Error in feature extraction: {str(e)}"
 
-# MIC prediction for AMP sequences
+# MIC prediction function
 def predictmic(sequence):
     sequence = ''.join([aa for aa in sequence.upper() if aa in "ACDEFGHIKLMNPQRSTVWY"])
     if len(sequence) < 10:
@@ -123,7 +98,7 @@ def predictmic(sequence):
 
     return mic_results
 
-# Full prediction pipeline
+# Main prediction function
 def full_prediction(sequence):
     features = extract_features(sequence)
     if isinstance(features, str):
@@ -131,9 +106,14 @@ def full_prediction(sequence):
 
     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[prediction] * 100, 2)
 
+    try:
+        class_index = list(model.classes_).index(prediction)
+        confidence = round(probabilities[class_index] * 100, 2)
+    except Exception:
+        confidence = "Unknown"
+
+    amp_result = "Antimicrobial Peptide (AMP)" if prediction == 0 else "Non-AMP"
     result = f"Prediction: {amp_result}\nConfidence: {confidence}%\n"
 
     if prediction == 0:
@@ -156,7 +136,7 @@ def full_prediction(sequence):
 
     return result
 
-# Gradio interface
+# Gradio UI
 iface = gr.Interface(
     fn=full_prediction,
     inputs=gr.Textbox(label="Enter Protein Sequence"),