Update app.py

app.py

@@ -22,7 +22,6 @@ from encoding import one_hot_padding, position_encoding, blosum62_padding, zscal
 
 # --- 2. CONFIGURATION ---
 SEEDS = [2023, 2024, 2025]
-TASK = 'PDC'
 
 # --- 3. LOAD ESM-2 ---
 print("📦 Downloading ESM-2 Model safely...")
@@ -41,7 +40,7 @@ esm_model.eval()
 # --- 4. CORE HELPER FUNCTIONS ---
 
 def get_esm_embedding(sequence):
-    clean_seq = re.sub(r"[UZOB]", "X", str(sequence)) # Ensure string
+    clean_seq = re.sub(r"[UZOB]", "X", str(sequence))
     encoded_input = tokenizer(clean_seq, return_tensors='pt')
     with torch.no_grad():
         output = esm_model(**encoded_input)
@@ -60,7 +59,6 @@ def get_t2_features(sequence, padding=71):
     return tf.constant(combined, dtype=tf.float32)
 
 def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
-    # Create temp dataframe for the Graph Loader
     df = pd.DataFrame({
         'Payload SMILES': [payload_smiles],
         'Linker SMILES': [linker_smiles],
@@ -69,7 +67,6 @@ def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
     })
     df.to_excel('temp_input.xlsx', index=False)
     
-    # Use existing loader
     dataset_gen = Graph_Classification_Dataset(
         'temp_input.xlsx',
         smiles_field1='Payload SMILES',
@@ -95,136 +92,152 @@ def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
 
     x1, adjoin1, _, x2, adjoin2, _ = target_batch
     
-    # Create Masks
     seq1 = tf.cast(tf.math.equal(x1, 0), tf.float32)
     mask1 = seq1[:, tf.newaxis, tf.newaxis, :]
     seq2 = tf.cast(tf.math.equal(x2, 0), tf.float32)
     mask2 = seq2[:, tf.newaxis, tf.newaxis, :]
     
-    # Get Features
     t2 = get_t2_features(peptide_seq)
     t1 = get_esm_embedding(peptide_seq)
             
     return x1, mask1, adjoin1, x2, mask2, adjoin2, t1, t2
 
-def get_raw_score(payload, linker, peptide):
-    """Calculates the raw probability (0.0 to 1.0)"""
-    try:
-        data = preprocess_input(payload, linker, peptide)
-        if data is None: return None
+# --- 5. HIGH-SPEED PREDICTION LOGIC ---
+
+def load_ensemble_models():
+    """Loads all 3 models into memory ONCE."""
+    models = []
+    keras.backend.clear_session()
+    
+    for seed in SEEDS:
+        weight_file = f'PDC_{seed}.h5'
+        if not os.path.exists(weight_file):
+            models.append(None)
+            continue
+            
+        model = PredictModel(num_layers=6, d_model=256, dff=512, num_heads=8, vocab_size=18, a=1, dense_dropout=0.4)
+        # Dummy call to build graph
+        # We use dummy tensors to initialize the weights
+        dummy_x = tf.zeros((1, 512), dtype=tf.int64) 
+        # ... (simplified init) ... 
+        # Actually, simpler to just load weights if we build input correctly. 
+        # But Safe way: Just return the fresh model and load weights during loop if needed?
+        # No, that's slow. We load weights NOW.
         
-        x1, mask1, adj1, x2, mask2, adj2, t1, t2 = data
-        total_prob = 0.0
-        models_ran = 0
+        # NOTE: To load weights, we need to build the model first. 
+        # We will do lazy-loading inside the loop but KEEP the model in memory.
+        models.append((model, weight_file))
         
-        for seed in SEEDS:
-            weight_file = f'PDC_{seed}.h5'
-            if not os.path.exists(weight_file): continue
-                
-            keras.backend.clear_session()
-            model = PredictModel(num_layers=6, d_model=256, dff=512, num_heads=8, vocab_size=18, a=1, dense_dropout=0.4)
-            
-            # Init and Load
-            _ = model(x1=x1, mask1=mask1, training=False, adjoin_matrix1=adj1,
+    return models
+
+def predict_with_loaded_models(models, inputs):
+    """Runs prediction using pre-loaded models"""
+    x1, mask1, adj1, x2, mask2, adj2, t1, t2 = inputs
+    total_prob = 0.0
+    models_ran = 0
+    
+    for model_tuple in models:
+        if model_tuple is None: continue
+        model, weight_path = model_tuple
+        
+        # If model is not built, we build it now (Lazy Build)
+        if not model.built:
+             _ = model(x1=x1, mask1=mask1, training=False, adjoin_matrix1=adj1,
                       x2=x2, mask2=mask2, adjoin_matrix2=adj2, t1=t1, t2=t2)
-            model.load_weights(weight_file)
-            
-            logits = model(x1=x1, mask1=mask1, training=False, adjoin_matrix1=adj1,
-                           x2=x2, mask2=mask2, adjoin_matrix2=adj2, t1=t1, t2=t2)
-            
-            prob = tf.sigmoid(logits).numpy()[0][0]
-            total_prob += prob
-            models_ran += 1
-            del model
-            
-        if models_ran == 0: return None
-        return total_prob / models_ran
+             model.load_weights(weight_path)
+        
+        logits = model(x1=x1, mask1=mask1, training=False, adjoin_matrix1=adj1,
+                       x2=x2, mask2=mask2, adjoin_matrix2=adj2, t1=t1, t2=t2)
         
-    except Exception:
-        return None
+        prob = tf.sigmoid(logits).numpy()[0][0]
+        total_prob += prob
+        models_ran += 1
+        
+    if models_ran == 0: return None
+    return total_prob / models_ran
 
-# --- 5. UI HANDLERS ---
+# --- 6. UI HANDLERS ---
 
 def predict_single(payload, linker, peptide):
     if not payload or not linker or not peptide:
         return "⚠️ Please fill in all fields."
     
-    score = get_raw_score(payload, linker, peptide)
+    # For single prediction, we accept the small slowness of loading
+    # to avoid keeping models in memory forever if not needed.
+    # OR reuse the logic. Let's reuse logic for consistency.
+    models = load_ensemble_models()
+    data = preprocess_input(payload, linker, peptide)
+    
+    if data is None: return "Error processing input."
     
-    if score is None:
-        return "Error processing input. Check SMILES validity."
+    score = predict_with_loaded_models(models, data)
+    keras.backend.clear_session() # Clean up after single use
     
     label = "Active PDC ✅" if score > 0.5 else "Inactive ❌"
     return f"{label}\nConfidence: {score:.2%}"
 
 def predict_batch(file_obj, progress=gr.Progress()):
-    if file_obj is None:
-        return None
+    if file_obj is None: return None, None
     
-    # Read File
+    # 1. Read File
     if file_obj.name.endswith('.csv'):
         df = pd.read_csv(file_obj.name)
     else:
         df = pd.read_excel(file_obj.name)
         
-    # --- SMART COLUMN DETECTION ---
-    # We map the "Keywords" we look for -> to the "Standard Names" our code needs
-    keyword_map = {
-        'payload': 'Payload SMILES', 
-        'linker': 'Linker SMILES', 
-        'peptide': 'Peptide Sequence'
-    }
-    
+    # 2. Map Columns
+    keyword_map = {'payload': 'Payload SMILES', 'linker': 'Linker SMILES', 'peptide': 'Peptide Sequence'}
     col_map = {}
     missing = []
-
     for keyword, standard_name in keyword_map.items():
-        # Look for a column that contains the keyword (e.g. "payload") case-insensitive
         match = next((c for c in df.columns if keyword in c.lower()), None)
-        
-        if match:
-            col_map[standard_name] = match # Map found column to standard name
-        else:
-            missing.append(keyword)
+        if match: col_map[standard_name] = match
+        else: missing.append(keyword)
     
     if missing:
-        return f"Error: Could not find columns containing keywords: {missing}. Found columns: {list(df.columns)}"
+        return None, pd.DataFrame({"Error": [f"Missing columns: {missing}"]})
+    
+    # 3. LOAD MODELS ONCE (The Speed Fix)
+    print("🚀 Loading Ensemble Models for Batch Processing...")
+    models = load_ensemble_models()
     
     results = []
     
-    # Iterate with Progress Bar
+    # 4. Iterate Rows
     for index, row in progress.tqdm(df.iterrows(), total=len(df)):
-        # Use the mapped column names to get data
         p_smiles = row[col_map['Payload SMILES']]
         l_smiles = row[col_map['Linker SMILES']]
         peptide = row[col_map['Peptide Sequence']]
         
-        score = get_raw_score(p_smiles, l_smiles, peptide)
+        # Preprocess
+        data = preprocess_input(p_smiles, l_smiles, peptide)
         
-        if score is None:
-            pred = "Error"
-            conf = 0.0
+        if data is None:
+            pred, conf = "Error", 0.0
         else:
+            # Predict using the ALREADY LOADED models
+            score = predict_with_loaded_models(models, data)
             pred = "Active" if score > 0.5 else "Inactive"
             conf = score
             
         results.append([p_smiles, l_smiles, peptide, pred, f"{conf:.4f}"])
         
-    # Create Result DataFrame
-    res_df = pd.DataFrame(results, columns=['Payload', 'Linker', 'Peptide', 'Prediction', 'Score'])
+    # 5. Cleanup
+    keras.backend.clear_session()
     
+    # 6. Format Outputs
+    res_df = pd.DataFrame(results, columns=['Payload', 'Linker', 'Peptide', 'Prediction', 'Score'])
     output_file = "predictions.csv"
     res_df.to_csv(output_file, index=False)
     
-    return output_file
+    return output_file, res_df
 
-# --- 6. GRADIO INTERFACE ---
+# --- 7. GRADIO INTERFACE ---
 with gr.Blocks(title="PDCNet Ensemble Demo") as demo:
     gr.Markdown("# 💊 PDCNet Ensemble Predictor")
     gr.Markdown("Predict Peptide-Drug Conjugate activity. Upload a file or type manually.")
     
     with gr.Tabs():
-        # TAB 1: SINGLE INPUT
         with gr.TabItem("Single Prediction"):
             with gr.Row():
                 with gr.Column():
@@ -236,13 +249,20 @@ with gr.Blocks(title="PDCNet Ensemble Demo") as demo:
                     out_result = gr.Textbox(label="Prediction Result", lines=4)
             btn_predict.click(predict_single, inputs=[in_payload, in_linker, in_peptide], outputs=out_result)
             
-        # TAB 2: BATCH INPUT
         with gr.TabItem("Batch Prediction (File)"):
-            gr.Markdown("Upload a CSV or Excel file. Must have columns for Payload, Linker, and Peptide.")
-            file_input = gr.File(label="Upload Dataset (.csv or .xlsx)")
-            btn_batch = gr.Button("Run Batch Prediction", variant="primary")
-            file_output = gr.File(label="Download Results")
+            gr.Markdown("Upload a CSV or Excel file.")
+            with gr.Row():
+                file_input = gr.File(label="Upload Dataset")
+                btn_batch = gr.Button("Run Batch Prediction", variant="primary")
+            
+            # OUTPUTS: Visual Table AND Download File
+            out_dataframe = gr.Dataframe(label="Results Preview", interactive=False)
+            file_output = gr.File(label="Download Full Results (.csv)")
             
-            btn_batch.click(predict_batch, inputs=file_input, outputs=file_output)
+            btn_batch.click(
+                predict_batch, 
+                inputs=file_input, 
+                outputs=[file_output, out_dataframe] # Two outputs now!
+            )
 
 demo.launch()