Update app.py

app.py

@@ -97,6 +97,7 @@ def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
     df.to_excel('temp_input.xlsx', index=False)
     
     # 2. Get Graph Features
+    # We rely on the existing loader, but we need to handle the random splitting
     dataset_gen = Graph_Classification_Dataset(
         'temp_input.xlsx',
         smiles_field1='Payload SMILES',
@@ -109,8 +110,25 @@ def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
         addH=True
     ).get_data()
     
-    _, _, val_dataset = dataset_gen
-    x1, adjoin1, _, x2, adjoin2, _ = next(iter(val_dataset))
+    # Unpack the three datasets (Train, Test, Val)
+    train_ds, test_ds, val_ds = dataset_gen
+    
+    # CRITICAL FIX: Check all 3 datasets for the data
+    # (Because with 1 row, it randomly lands in one and leaves the others empty)
+    target_batch = None
+    for ds in [train_ds, test_ds, val_ds]:
+        try:
+            # Try to grab the batch
+            target_batch = next(iter(ds))
+            if target_batch is not None:
+                break # Found it!
+        except Exception:
+            continue # This one was empty, check the next
+            
+    if target_batch is None:
+        return None, None, None, None, None, None, None, None
+
+    x1, adjoin1, _, x2, adjoin2, _ = target_batch
     
     # Masks
     seq1 = tf.cast(tf.math.equal(x1, 0), tf.float32)
@@ -121,7 +139,7 @@ def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
     # 3. Get Peptide Features
     t2 = get_t2_features(peptide_seq)
     
-    # DYNAMIC GENERATION (The Fix)
+    # DYNAMIC GENERATION
     t1 = get_protbert_embedding(peptide_seq)
             
     return x1, mask1, adjoin1, x2, mask2, adjoin2, t1, t2
@@ -173,7 +191,7 @@ import re # Needed for regex in protbert cleaning
 
 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 ProtBERT embedding.")
+    gr.Markdown("Predict the activity of Peptide-Drug Conjugates.")
     
     with gr.Row():
         with gr.Column():