Update app.py

app.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 # app.py - RLAnOxPeptide Gradio Web Application
-# Final version updated to use an AdvancedProtT5Generator with a LoRA backbone.
+
 
 import os
 import torch
@@ -15,6 +15,7 @@ from sklearn.cluster import KMeans
 from tqdm import tqdm
 import transformers
 import time
+import copy  # ✅ ADDED: For deep copying the base model
 
 # NEW DEPENDENCY: peft library for LoRA
 from peft import PeftModel
@@ -34,20 +35,21 @@ token2id["<EOS>"] = 1
 id2token = {i: t for t, i in token2id.items()}
 VOCAB_SIZE = len(token2id)
 
+
 # --- Validator's Feature Extractor Class ---
+# ✅ MODIFIED: Accepts a pre-loaded model instead of loading its own.
 class LoRAProtT5Extractor:
-    def __init__(self, base_model_id, lora_adapter_path):
+    def __init__(self, preloaded_base_model, preloaded_tokenizer, lora_adapter_path):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"Initializing Validator Feature Extractor on device: {self.device}")
 
-        print(f"  - [Validator] Loading base model and tokenizer from '{base_model_id}'...")
-        base_model = transformers.T5EncoderModel.from_pretrained(base_model_id)
-        self.tokenizer = transformers.T5Tokenizer.from_pretrained(base_model_id)
+        base_model = preloaded_base_model
+        self.tokenizer = preloaded_tokenizer
         
         if not os.path.exists(lora_adapter_path):
             raise FileNotFoundError(f"Error: Validator LoRA adapter directory not found at: {lora_adapter_path}")
             
-        print(f"  - [Validator] Loading and applying LoRA adapter from: {lora_adapter_path}")
+        print(f"  - [Validator] Applying LoRA adapter from: {lora_adapter_path}")
         lora_model = PeftModel.from_pretrained(base_model, lora_adapter_path)
         
         print("  - [Validator] Merging LoRA weights for faster inference...")
@@ -58,14 +60,11 @@ class LoRAProtT5Extractor:
     def encode(self, sequence):
         if not sequence or not isinstance(sequence, str):
             return np.zeros((1, 1024), dtype=np.float32)
-
         seq_spaced = " ".join(list(sequence))
         encoded_input = self.tokenizer(seq_spaced, return_tensors='pt', padding=True, truncation=True)
         encoded_input = {k: v.to(self.device) for k, v in encoded_input.items()}
-
         with torch.no_grad():
             embedding = self.model(**encoded_input).last_hidden_state
-        
         emb_np = embedding.squeeze(0).cpu().numpy()
         return emb_np if emb_np.shape[0] > 0 else np.zeros((1, 1024), dtype=np.float32)
 
@@ -78,15 +77,12 @@ class AntioxidantPredictor(nn.Module):
         self.handcrafted_dim = input_dim - self.prott5_dim
         self.seq_len = 16
         self.prott5_feature_dim = 64
-
         encoder_layer = nn.TransformerEncoderLayer(d_model=self.prott5_feature_dim, nhead=transformer_heads, dropout=transformer_dropout, batch_first=True)
         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=transformer_layers)
-
         fused_dim = self.prott5_feature_dim + self.handcrafted_dim
         self.fusion_fc = nn.Sequential(nn.Linear(fused_dim, 1024), nn.ReLU(), nn.Dropout(0.3), nn.Linear(1024, 512), nn.ReLU(), nn.Dropout(0.3))
         self.classifier = nn.Sequential(nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, 1))
         self.temperature = nn.Parameter(torch.ones(1), requires_grad=False)
-
     def forward(self, x):
         batch_size = x.size(0)
         prot_t5_features = x[:, :self.prott5_dim]
@@ -98,26 +94,23 @@ class AntioxidantPredictor(nn.Module):
         fused_output = self.fusion_fc(fused_features)
         logits = self.classifier(fused_output)
         return logits / self.temperature
-
     def get_temperature(self):
         return self.temperature.item()
 
 
-# --- ✅ NEW Generator Model Architecture ---
+# --- Generator Model Architecture ---
+# ✅ MODIFIED: Accepts a pre-loaded model instead of loading its own.
 class AdvancedProtT5Generator(nn.Module):
-    def __init__(self, base_model_id, lora_adapter_path, vocab_size):
+    def __init__(self, preloaded_base_model, lora_adapter_path, vocab_size):
         super(AdvancedProtT5Generator, self).__init__()
         
-        print(f"  - [Generator] Loading base ProtT5 model from '{base_model_id}'...")
-        base_model = transformers.T5EncoderModel.from_pretrained(base_model_id)
-
+        base_model = preloaded_base_model
         print(f"  - [Generator] Applying LoRA adapter from: {lora_adapter_path}")
         self.backbone = PeftModel.from_pretrained(base_model, lora_adapter_path)
         
-        # Expose the embedding layer for the clustering function
         self.embed_tokens = self.backbone.get_input_embeddings()
         
-        embed_dim = self.backbone.config.d_model # Should be 1024
+        embed_dim = self.backbone.config.d_model
         self.lm_head = nn.Linear(embed_dim, vocab_size)
         
         self.vocab_size = vocab_size
@@ -164,7 +157,7 @@ try:
 except ImportError:
     raise gr.Error("Fatal Error: `feature_extract.py` not found. This file is required. Please upload it to your repository.")
 
-# --- ✅ UPDATED Clustering Logic ---
+# --- Clustering Logic (Unchanged) ---
 def cluster_sequences(generator, sequences, num_clusters, device):
     if not sequences or len(sequences) < num_clusters:
         return sequences[:num_clusters]
@@ -176,14 +169,10 @@ def cluster_sequences(generator, sequences, num_clusters, device):
             ids += [token2id["<PAD>"]] * (max_len - len(ids))
             token_ids_list.append(ids)
         input_ids = torch.tensor(token_ids_list, dtype=torch.long, device=device)
-        
-        # Use the generator's exposed embedding layer
         embeddings = generator.embed_tokens(input_ids)
-
         mask = (input_ids != token2id["<PAD>"]).unsqueeze(-1).float()
         seq_embeds = (embeddings * mask).sum(dim=1) / (mask.sum(dim=1) + 1e-9)
         seq_embeds_np = seq_embeds.cpu().numpy()
-        
     kmeans = KMeans(n_clusters=int(num_clusters), random_state=42, n_init='auto').fit(seq_embeds_np)
     reps = []
     for i in range(int(num_clusters)):
@@ -205,21 +194,25 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 try:
     # --- Define file paths ---
     PROTT5_BASE_MODEL_ID = "Rostlab/prot_t5_xl_uniref50"
-    
-    # Paths for the Validator System
     VALIDATOR_LORA_PATH = "./lora_finetuned_prott5"
     PREDICTOR_HEAD_CHECKPOINT_PATH = "./predictor_with_lora_checkpoints/final_predictor_with_lora.pth"
     SCALER_PATH = "./predictor_with_lora_checkpoints/scaler_lora.pkl"
-
-    # Paths for the Generator System
     GENERATOR_LORA_DIR = "./generator_with_lora_output/final_lora_generator"
     GENERATOR_LM_HEAD_PATH = os.path.join(GENERATOR_LORA_DIR, "lm_head.pth")
 
+    # ✅ OPTIMIZED: Load the base model and tokenizer only ONCE
+    print(f"--- Loading Base ProtT5 Model ({PROTT5_BASE_MODEL_ID}) just once... ---")
+    base_prot_t5_model = transformers.T5EncoderModel.from_pretrained(PROTT5_BASE_MODEL_ID)
+    base_tokenizer = transformers.T5Tokenizer.from_pretrained(PROTT5_BASE_MODEL_ID)
+    print("✅ Base ProtT5 Model loaded.")
+
     # --- Load Validator System ---
-    print("--- Loading Validator System ---")
+    print("\n--- Initializing Validator System ---")
     VALIDATOR_SCALER = joblib.load(SCALER_PATH)
+    # Pass a deep copy of the base model to prevent modification conflicts
     VALIDATOR_EXTRACTOR = LoRAProtT5Extractor(
-        base_model_id=PROTT5_BASE_MODEL_ID,
+        preloaded_base_model=copy.deepcopy(base_prot_t5_model),
+        preloaded_tokenizer=base_tokenizer,
         lora_adapter_path=VALIDATOR_LORA_PATH
     )
     PREDICTOR_MODEL = AntioxidantPredictor(input_dim=1914)
@@ -229,9 +222,10 @@ try:
     print("✅ Validator System loaded successfully.")
 
     # --- Load Generator System ---
-    print("\n--- Loading Generator System ---")
+    print("\n--- Initializing Generator System ---")
+    # Pass a deep copy of the base model here as well
     GENERATOR_MODEL = AdvancedProtT5Generator(
-        base_model_id=PROTT5_BASE_MODEL_ID,
+        preloaded_base_model=copy.deepcopy(base_prot_t5_model),
         lora_adapter_path=GENERATOR_LORA_DIR,
         vocab_size=VOCAB_SIZE
     )
@@ -254,10 +248,9 @@ except Exception as e:
 
 def predict_peptide_wrapper(sequence_str):
     if not sequence_str or not isinstance(sequence_str, str) or any(c not in AMINO_ACIDS for c in sequence_str.upper()):
-        return "0.0000", "Error: Please enter a valid peptide sequence using standard amino acids (ACDEFGHIKLMNPQRSTVWY)."
+        return "0.0000", "Error: Please enter a valid peptide sequence using standard amino acids."
     
     try:
-        # Use the VALIDATOR's feature extractor
         features = extract_features(sequence_str.upper(), VALIDATOR_EXTRACTOR, L_fixed=29, d_model_pe=16)
         scaled_features = VALIDATOR_SCALER.transform(features.reshape(1, -1))