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("""

💊 PDCNet Ensemble Predictor

Predict Peptide-Drug Conjugate activity using Graph Neural Networks + ESM-2 Protein Language Model.

""") 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)