Update app.py

app.py

@@ -12,7 +12,7 @@ import tf_keras as keras
 import numpy as np
 import pandas as pd
 import torch
-from transformers import AutoTokenizer, EsmModel # CHANGED: Switched to ESM
+from transformers import AutoTokenizer, EsmModel
 from huggingface_hub import snapshot_download
 from model import PredictModel
 
@@ -24,10 +24,8 @@ from encoding import one_hot_padding, position_encoding, blosum62_padding, zscal
 SEEDS = [2023, 2024, 2025]
 TASK = 'PDC'
 
-# --- 3. LOAD ESM-2 (The Correct "Brain" for your weights) ---
+# --- 3. LOAD ESM-2 ---
 print("📦 Downloading ESM-2 Model safely...")
-
-# We use the 150M parameter model which has dim=640 (Matches your 1664 input size)
 REPO_ID = "facebook/esm2_t30_150M_UR50D"
 
 model_path = snapshot_download(
@@ -38,43 +36,21 @@ model_path = snapshot_download(
 
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 esm_model = EsmModel.from_pretrained(model_path)
-
-# Move to CPU
 esm_model.eval()
 
-# --- 4. HELPER FUNCTIONS ---
+# --- 4. CORE HELPER FUNCTIONS ---
 
 def get_esm_embedding(sequence):
-    """
-    Generates the t1 tensor dynamically using ESM-2 (640 dim).
-    """
-    # ESM does NOT need spaces, but handles standard strings
-    # Clean sequence
-    clean_seq = re.sub(r"[UZOB]", "X", sequence)
-
-    # Tokenize
+    clean_seq = re.sub(r"[UZOB]", "X", str(sequence)) # Ensure string
     encoded_input = tokenizer(clean_seq, return_tensors='pt')
-    
-    # Run Model
     with torch.no_grad():
         output = esm_model(**encoded_input)
-        
-    # Get the Last Hidden State [1, Seq_Len, 640]
-    embedding = output.last_hidden_state
-    
-    # Remove [CLS] and [SEP] (first and last tokens)
-    embedding = embedding[:, 1:-1, :] 
-    
-    # Mean Pooling: [1, Seq_Len, 640] -> [1, 640]
+    embedding = output.last_hidden_state[:, 1:-1, :] 
     pooled_embedding = torch.mean(embedding, dim=1)
-    
-    # Convert to TensorFlow
-    tf_embedding = tf.convert_to_tensor(pooled_embedding.numpy(), dtype=tf.float32)
-    
-    return tf_embedding
+    return tf.convert_to_tensor(pooled_embedding.numpy(), dtype=tf.float32)
 
 def get_t2_features(sequence, padding=71):
-    selected_sequences = [sequence]
+    selected_sequences = [str(sequence)]
     one_hot = one_hot_padding(selected_sequences, padding)
     pos = position_encoding(padding)
     blosum = blosum62_padding(selected_sequences, padding)
@@ -84,7 +60,7 @@ def get_t2_features(sequence, padding=71):
     return tf.constant(combined, dtype=tf.float32)
 
 def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
-    # 1. Create dummy input
+    # Create temp dataframe for the Graph Loader
     df = pd.DataFrame({
         'Payload SMILES': [payload_smiles],
         'Linker SMILES': [linker_smiles],
@@ -93,7 +69,7 @@ def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
     })
     df.to_excel('temp_input.xlsx', index=False)
     
-    # 2. Get Graph Features
+    # Use existing loader
     dataset_gen = Graph_Classification_Dataset(
         'temp_input.xlsx',
         smiles_field1='Payload SMILES',
@@ -106,7 +82,7 @@ def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
         addH=True
     ).get_data()
     
-    # Robust Batch Retrieval (Check Train/Test/Val)
+    # Robust Batch Retrieval
     train_ds, test_ds, val_ds = dataset_gen
     target_batch = None
     for ds in [train_ds, test_ds, val_ds]:
@@ -115,49 +91,42 @@ def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
             if target_batch is not None: break
         except: continue
             
-    if target_batch is None:
-        return None, None, None, None, None, None, None, None
+    if target_batch is None: return None
 
     x1, adjoin1, _, x2, adjoin2, _ = target_batch
     
-    # Masks
+    # 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, :]
     
-    # 3. Get Features
+    # Get Features
     t2 = get_t2_features(peptide_seq)
-    t1 = get_esm_embedding(peptide_seq) # Use New ESM Function
+    t1 = get_esm_embedding(peptide_seq)
             
     return x1, mask1, adjoin1, x2, mask2, adjoin2, t1, t2
 
-def predict_ensemble(payload, linker, peptide):
-    if not payload or not linker or not peptide:
-        return "⚠️ Please fill in all fields."
-        
+def get_raw_score(payload, linker, peptide):
+    """Calculates the raw probability (0.0 to 1.0)"""
     try:
-        x1, mask1, adj1, x2, mask2, adj2, t1, t2 = preprocess_input(payload, linker, peptide)
+        data = preprocess_input(payload, linker, peptide)
+        if data is None: return None
         
-        if x1 is None:
-            return "Error: Could not process chemical structures. Check SMILES validity."
-
+        x1, mask1, adj1, x2, mask2, adj2, t1, t2 = data
         total_prob = 0.0
         models_ran = 0
         
         for seed in SEEDS:
             weight_file = f'PDC_{seed}.h5'
-            if not os.path.exists(weight_file):
-                continue
+            if not os.path.exists(weight_file): continue
                 
             keras.backend.clear_session()
-            # Re-init model
             model = PredictModel(num_layers=6, d_model=256, dff=512, num_heads=8, vocab_size=18, a=1, dense_dropout=0.4)
             
-            # Dummy call
+            # Init and Load
             _ = 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,
@@ -168,31 +137,112 @@ def predict_ensemble(payload, linker, peptide):
             models_ran += 1
             del model
             
-        if models_ran == 0:
-            return "Error: No model weights found on server."
+        if models_ran == 0: return None
+        return total_prob / models_ran
+        
+    except Exception:
+        return None
+
+# --- 5. 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)
+    
+    if score is None:
+        return "Error processing input. Check SMILES validity."
+    
+    label = "Active PDC ✅" if score > 0.5 else "Inactive ❌"
+    return f"{label}\nConfidence: {score:.2%}"
 
-        avg_score = total_prob / models_ran
-        label = "Active PDC ✅" if avg_score > 0.5 else "Inactive ❌"
-        return f"{label}\nConfidence: {avg_score:.2%}"
+def predict_batch(file_obj, progress=gr.Progress()):
+    if file_obj is None:
+        return None
+    
+    # Read File
+    if file_obj.name.endswith('.csv'):
+        df = pd.read_csv(file_obj.name)
+    else:
+        df = pd.read_excel(file_obj.name)
         
-    except Exception as e:
-        return f"Error: {str(e)}"
+    # --- 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'
+    }
+    
+    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 missing:
+        return f"Error: Could not find columns containing keywords: {missing}. Found columns: {list(df.columns)}"
+    
+    results = []
+    
+    # Iterate with Progress Bar
+    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)
+        
+        if score is None:
+            pred = "Error"
+            conf = 0.0
+        else:
+            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'])
+    
+    output_file = "predictions.csv"
+    res_df.to_csv(output_file, index=False)
+    
+    return output_file
 
-# --- 5. GRADIO INTERFACE ---
+# --- 6. GRADIO INTERFACE ---
 with gr.Blocks(title="PDCNet Ensemble Demo") as demo:
     gr.Markdown("# 💊 PDCNet Ensemble Predictor")
-    gr.Markdown("Predict the activity of Peptide-Drug Conjugates. Supports **Any Peptide Sequence** via live ESM-2 embedding.")
+    gr.Markdown("Predict Peptide-Drug Conjugate activity. Upload a file or type manually.")
     
-    with gr.Row():
-        with gr.Column():
-            in_payload = gr.Textbox(label="Payload SMILES", placeholder="e.g. COc1cc2c...", lines=2)
-            in_linker = gr.Textbox(label="Linker SMILES", placeholder="e.g. C(=O)CCCC(=O)O", lines=2)
-            in_peptide = gr.Textbox(label="Peptide Sequence", placeholder="e.g. ACDEFGH", lines=2)
-            btn_predict = gr.Button("Predict Activity", variant="primary")
+    with gr.Tabs():
+        # TAB 1: SINGLE INPUT
+        with gr.TabItem("Single Prediction"):
+            with gr.Row():
+                with gr.Column():
+                    in_payload = gr.Textbox(label="Payload SMILES", placeholder="e.g. COc1cc2c...", lines=2)
+                    in_linker = gr.Textbox(label="Linker SMILES", placeholder="e.g. C(=O)CCCC(=O)O", lines=2)
+                    in_peptide = gr.Textbox(label="Peptide Sequence", placeholder="e.g. ACDEFGH", lines=2)
+                    btn_predict = gr.Button("Predict Activity", variant="primary")
+                with gr.Column():
+                    out_result = gr.Textbox(label="Prediction Result", lines=4)
+            btn_predict.click(predict_single, inputs=[in_payload, in_linker, in_peptide], outputs=out_result)
             
-        with gr.Column():
-            out_result = gr.Textbox(label="Prediction Result", lines=4)
+        # 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")
             
-    btn_predict.click(predict_ensemble, inputs=[in_payload, in_linker, in_peptide], outputs=out_result)
+            btn_batch.click(predict_batch, inputs=file_input, outputs=file_output)
 
 demo.launch()