Create app.py

app.py

@@ -0,0 +1,181 @@
+import os
+
+# --- 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 pickle
+import torch
+from transformers import BertModel, BertTokenizer
+from model import PredictModel
+
+# Import your local helper files
+# (These must be uploaded to the same folder as app.py)
+from dataset import Graph_Classification_Dataset
+from encoding import one_hot_padding, position_encoding, blosum62_padding, zscale_padding
+
+# --- 2. CONFIGURATION ---
+SEEDS = [2023, 2024, 2025]
+TASK = 'PDC'
+
+# --- 3. LOAD PROTBERT (The "Brain" for Peptides) ---
+print("📦 Downloading/Loading ProtBERT Model... (This runs once)")
+tokenizer = BertTokenizer.from_pretrained("Rostlab/prot_bert", do_lower_case=False)
+bert_model = BertModel.from_pretrained("Rostlab/prot_bert")
+
+# Move to CPU (Hugging Face Free Tier is CPU-only usually)
+bert_model.eval()
+
+# --- 4. HELPER FUNCTIONS ---
+
+def get_protbert_embedding(sequence):
+    """
+    Generates the t1 tensor dynamically for ANY peptide sequence using ProtBERT.
+    Replaces Peptide.pkl lookup.
+    """
+    # ProtBERT expects spaces between amino acids: "A C D E"
+    spaced_sequence = " ".join(list(sequence))
+    
+    # Clean sequence (remove weird chars if any)
+    spaced_sequence = re.sub(r"[UZOB]", "X", spaced_sequence)
+
+    # Tokenize
+    encoded_input = tokenizer(spaced_sequence, return_tensors='pt')
+    
+    # Run Model (No Gradient needed for inference)
+    with torch.no_grad():
+        output = bert_model(**encoded_input)
+        
+    # Get the Last Hidden State [1, Seq_Len, 1024]
+    # We remove the [CLS] and [SEP] tokens (first and last)
+    embedding = output.last_hidden_state
+    embedding = embedding[:, 1:-1, :] 
+    
+    # Convert to TensorFlow Tensor
+    tf_embedding = tf.convert_to_tensor(embedding.numpy(), dtype=tf.float32)
+    
+    # Ensure shape matches [1, Seq_Len, 1024]
+    return tf_embedding
+
+def get_t2_features(sequence, padding=71):
+    """Replicates the logic from extract_t2 in class.py"""
+    selected_sequences = [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):
+    # 1. Create dummy input for Graph Dataset
+    df = pd.DataFrame({
+        'Payload SMILES': [payload_smiles],
+        'Linker SMILES': [linker_smiles],
+        'label': [0],
+        'index': [0]
+    })
+    df.to_excel('temp_input.xlsx', index=False)
+    
+    # 2. Get Graph Features
+    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()
+    
+    _, _, val_dataset = dataset_gen
+    x1, adjoin1, _, x2, adjoin2, _ = next(iter(val_dataset))
+    
+    # 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 Peptide Features
+    t2 = get_t2_features(peptide_seq)
+    
+    # DYNAMIC GENERATION (The Fix)
+    t1 = get_protbert_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."
+        
+    try:
+        x1, mask1, adj1, x2, mask2, adj2, t1, t2 = preprocess_input(payload, linker, peptide)
+        
+        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
+                
+            keras.backend.clear_session()
+            model = PredictModel(num_layers=6, d_model=256, dff=512, num_heads=8, vocab_size=18, a=1, dense_dropout=0.4)
+            
+            # Init weights
+            _ = 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,
+                           x2=x2, mask2=mask2, adjoin_matrix2=adj2, t1=t1, t2=t2)
+            
+            prob = tf.sigmoid(logits).numpy()[0][0]
+            total_prob += prob
+            models_ran += 1
+            del model
+            
+        if models_ran == 0:
+            return "Error: No model weights found on server."
+
+        avg_score = total_prob / models_ran
+        label = "Active PDC ✅" if avg_score > 0.5 else "Inactive ❌"
+        return f"{label}\nConfidence: {avg_score:.2%}"
+        
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# --- 5. GRADIO INTERFACE ---
+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.")
+    
+    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_ensemble, inputs=[in_payload, in_linker, in_peptide], outputs=out_result)
+
+demo.launch()