Update app.py

app.py

@@ -117,15 +117,7 @@ def load_ensemble_models():
             
         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.
-        
-        # 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))
         
     return models
@@ -162,16 +154,13 @@ def predict_single(payload, linker, peptide):
     if not payload or not linker or not peptide:
         return "⚠️ Please fill in all fields."
     
-    # 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."
     
     score = predict_with_loaded_models(models, data)
-    keras.backend.clear_session() # Clean up after single use
+    keras.backend.clear_session()
     
     label = "Active PDC ✅" if score > 0.5 else "Inactive ❌"
     return f"{label}\nConfidence: {score:.2%}"
@@ -197,7 +186,7 @@ def predict_batch(file_obj, progress=gr.Progress()):
     if missing:
         return None, pd.DataFrame({"Error": [f"Missing columns: {missing}"]})
     
-    # 3. LOAD MODELS ONCE (The Speed Fix)
+    # 3. LOAD MODELS ONCE
     print("🚀 Loading Ensemble Models for Batch Processing...")
     models = load_ensemble_models()
     
@@ -209,36 +198,52 @@ def predict_batch(file_obj, progress=gr.Progress()):
         l_smiles = row[col_map['Linker SMILES']]
         peptide = row[col_map['Peptide Sequence']]
         
-        # Preprocess
         data = preprocess_input(p_smiles, l_smiles, peptide)
         
         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}"])
+        # UI UPDATE: Put Prediction and Score FIRST
+        results.append([pred, f"{conf:.4f}", p_smiles, l_smiles, peptide])
         
     # 5. Cleanup
     keras.backend.clear_session()
     
     # 6. Format Outputs
-    res_df = pd.DataFrame(results, columns=['Payload', 'Linker', 'Peptide', 'Prediction', 'Score'])
+    # Columns Reordered here for both display and CSV
+    res_df = pd.DataFrame(results, columns=['Prediction', 'Score', 'Payload', 'Linker', 'Peptide'])
     output_file = "predictions.csv"
     res_df.to_csv(output_file, index=False)
     
     return output_file, res_df
 
 # --- 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.")
+# CSS to center the content nicely
+css = """
+.gradio-container {text-align: center;}
+h1 {text-align: center; display: block;}
+p {text-align: center; display: block;}
+"""
+
+with gr.Blocks(title="PDCNet Ensemble Demo", css=css) as demo:
+    
+    # CENTERED HEADER
+    gr.HTML("""
+    <div style="text-align: center; max-width: 800px; margin: 0 auto;">
+        <h1>💊 PDCNet Ensemble Predictor</h1>
+        <p style="font-size: 1.1em;">
+            Predict Peptide-Drug Conjugate activity using Graph Neural Networks + ESM-2 Protein Language Model.
+        </p>
+    </div>
+    """)
     
     with gr.Tabs():
         with gr.TabItem("Single Prediction"):
+            # Centering the input area slightly
             with gr.Row():
                 with gr.Column():
                     in_payload = gr.Textbox(label="Payload SMILES", placeholder="e.g. COc1cc2c...", lines=2)
@@ -250,19 +255,20 @@ with gr.Blocks(title="PDCNet Ensemble Demo") as demo:
             btn_predict.click(predict_single, inputs=[in_payload, in_linker, in_peptide], outputs=out_result)
             
         with gr.TabItem("Batch Prediction (File)"):
-            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")
+            gr.Markdown("### Upload a CSV or Excel file to process many molecules at once.")
+            
+            # UI UPDATE: Stacked vertically (No gr.Row)
+            file_input = gr.File(label="Upload Dataset")
+            btn_batch = gr.Button("Run Batch Prediction", variant="primary")
             
-            # OUTPUTS: Visual Table AND Download File
+            # OUTPUTS
             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, out_dataframe] # Two outputs now!
+                outputs=[file_output, out_dataframe]
             )
 
 demo.launch()