Blame for app.py

import os
import re

# --- 1. CRITICAL SETUP FOR LEGACY KERAS ---
os.environ["TF_USE_LEGACY_KERAS"] = "1"
os.environ["TF_XLA_FLAGS"] = "--tf_xla_enable_xla_devices=false"
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"

import gradio as gr
import tensorflow as tf
import tf_keras as keras
import numpy as np
import pandas as pd
import torch
from transformers import AutoTokenizer, EsmModel
from huggingface_hub import snapshot_download
from model import PredictModel

# Import your local helper files
from dataset import Graph_Classification_Dataset
from encoding import one_hot_padding, position_encoding, blosum62_padding, zscale_padding

# --- 2. CONFIGURATION ---
SEEDS = [2023, 2024, 2025]

# --- 3. LOAD ESM-2 ---
print("📦 Downloading ESM-2 Model safely...")
REPO_ID = "facebook/esm2_t30_150M_UR50D"

model_path = snapshot_download(
    repo_id=REPO_ID,
    allow_patterns=["*.json", "pytorch_model.bin", "vocab.txt"], 
    local_files_only=False
)

tokenizer = AutoTokenizer.from_pretrained(model_path)
esm_model = EsmModel.from_pretrained(model_path)
esm_model.eval()

# --- 4. CORE HELPER FUNCTIONS ---

def get_esm_embedding(sequence):
    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)
    embedding = output.last_hidden_state[:, 1:-1, :] 
    pooled_embedding = torch.mean(embedding, dim=1)
    return tf.convert_to_tensor(pooled_embedding.numpy(), dtype=tf.float32)

def get_t2_features(sequence, padding=71):
    selected_sequences = [str(sequence)]
    one_hot = one_hot_padding(selected_sequences, padding)
    pos = position_encoding(padding)
    blosum = blosum62_padding(selected_sequences, padding)
    zscale = zscale_padding(selected_sequences, padding)
    pos_expanded = np.expand_dims(pos, axis=0).repeat(len(selected_sequences), axis=0)
    combined = np.concatenate((one_hot, pos_expanded, blosum, zscale), axis=2)
    return tf.constant(combined, dtype=tf.float32)

def preprocess_input(payload_smiles, linker_smiles, peptide_seq):
    df = pd.DataFrame({
        'Payload SMILES': [payload_smiles],
        'Linker SMILES': [linker_smiles],
        'label': [0],
        'index': [0]
    })
    df.to_excel('temp_input.xlsx', index=False)
    
    dataset_gen = Graph_Classification_Dataset(
        'temp_input.xlsx',
        smiles_field1='Payload SMILES',
        smiles_field2='Linker SMILES',
        label_field=['label'],
        index_field=['index'],
        seed=2023,
        batch_size=1,
        a=1,
        addH=True
    ).get_data()
    
    # Robust Batch Retrieval
    train_ds, test_ds, val_ds = dataset_gen
    target_batch = None
    for ds in [train_ds, test_ds, val_ds]:
        try:
            target_batch = next(iter(ds))
            if target_batch is not None: break
        except: continue
            
    if target_batch is None: return None

    x1, adjoin1, _, x2, adjoin2, _ = target_batch
    
    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, :]
    
    t2 = get_t2_features(peptide_seq)
    t1 = get_esm_embedding(peptide_seq)
            
    return x1, mask1, adjoin1, x2, mask2, adjoin2, t1, t2

# --- 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
        dummy_x = tf.zeros((1, 512), dtype=tf.int64) 
        models.append((model, weight_file))
        
    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_path)
        
        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
        
    if models_ran == 0: return None
    return total_prob / models_ran

# --- 6. UI HANDLERS ---

def predict_single(payload, linker, peptide):
    if not payload or not linker or not peptide:
        return "⚠️ Please fill in all fields."
    
    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()
    
    label = "Active PDC ✅" if score > 0.5 else "Inactive ❌"
    return f"{label}\nConfidence: {score:.3%}"

def predict_batch(file_obj, progress=gr.Progress()):
    if file_obj is None: return None, None
    
    # 1. Read File
    if file_obj.name.endswith('.csv'):
        df = pd.read_csv(file_obj.name)
    else:
        df = pd.read_excel(file_obj.name)
        
    # 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():
        match = next((c for c in df.columns if keyword in c.lower()), None)
        if match: col_map[standard_name] = match
        else: missing.append(keyword)
    
    if missing:
        return None, pd.DataFrame({"Error": [f"Missing columns: {missing}"]})
    
    # 3. LOAD MODELS ONCE
    print("🚀 Loading Ensemble Models for Batch Processing...")
    models = load_ensemble_models()
    
    results = []
    
    # 4. Iterate Rows
    for index, row in progress.tqdm(df.iterrows(), total=len(df)):
        p_smiles = row[col_map['Payload SMILES']]
        l_smiles = row[col_map['Linker SMILES']]
        peptide = row[col_map['Peptide Sequence']]
        
        data = preprocess_input(p_smiles, l_smiles, peptide)
        
        if data is None:
            pred, conf = "Error", 0.0
        else:
            score = predict_with_loaded_models(models, data)
            pred = "Active ✅" if score > 0.5 else "Inactive ❌"
            conf = score
            
        # UI UPDATE: Put Prediction and Score FIRST
        results.append([pred, f"{conf:.3%}", p_smiles, l_smiles, peptide])
        
    # 5. Cleanup
    keras.backend.clear_session()
    
    # 6. Format Outputs
    # 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 ---
# 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)
                    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.TabItem("Batch Prediction (File)"):
            gr.Markdown("### Upload a CSV or Excel file to process many molecules at once. Must have Payload, Linker and Peptide columns.")
            
            # UI UPDATE: Stacked vertically (No gr.Row)
            file_input = gr.File(label="Upload Dataset")
            btn_batch = gr.Button("Run Batch Prediction", variant="primary")
            
            # OUTPUTS
            out_dataframe = gr.Dataframe(
                label="Results Preview", 
                interactive=False, 
                wrap=True,  # <--- CRITICAL: Allows long SMILES strings to break lines
                column_widths=["150px", "150px", "minmax(200px, 1fr)", "minmax(150px, 1fr)", "minmax(150px, 1fr)"] 
            )
            file_output = gr.File(label="Download Full Results (.csv)")
            
            btn_batch.click(
                predict_batch, 
                inputs=file_input, 
                outputs=[file_output, out_dataframe]
            )

demo.launch(auth=(os.environ.get("app_user"), os.environ.get("app_pass")),
    ssr_mode=False)