Update app.py

app.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 # app.py - RLAnOxPeptide Gradio Web Application
-# Final version updated to use a LoRA-finetuned model for feature extraction.
+# Final version updated to use an AdvancedProtT5Generator with a LoRA backbone.
 
 import os
 import torch
@@ -34,28 +34,26 @@ token2id["<EOS>"] = 1
 id2token = {i: t for t, i in token2id.items()}
 VOCAB_SIZE = len(token2id)
 
-
-# --- LoRA Feature Extractor Model Class ---
-# ✅ REPLACED: This new class handles loading the base model and attaching the LoRA adapter.
+# --- Validator's Feature Extractor Class ---
 class LoRAProtT5Extractor:
     def __init__(self, base_model_id, lora_adapter_path):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Initializing LoRA-enhanced ProtT5 on device: {self.device}")
+        print(f"Initializing Validator Feature Extractor on device: {self.device}")
 
-        print(f"  - Loading base model and tokenizer from '{base_model_id}'...")
+        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)
         
         if not os.path.exists(lora_adapter_path):
-            raise FileNotFoundError(f"Error: LoRA adapter directory not found at: {lora_adapter_path}")
+            raise FileNotFoundError(f"Error: Validator LoRA adapter directory not found at: {lora_adapter_path}")
             
-        print(f"  - Loading and applying LoRA adapter from: {lora_adapter_path}")
+        print(f"  - [Validator] Loading and applying LoRA adapter from: {lora_adapter_path}")
         lora_model = PeftModel.from_pretrained(base_model, lora_adapter_path)
         
-        print("  - Merging LoRA weights for faster inference...")
+        print("  - [Validator] Merging LoRA weights for faster inference...")
         self.model = lora_model.merge_and_unload().to(self.device)
         self.model.eval()
-        print("  - LoRA-enhanced feature extractor is ready.")
+        print("  - Validator feature extractor is ready.")
 
     def encode(self, sequence):
         if not sequence or not isinstance(sequence, str):
@@ -71,7 +69,8 @@ class LoRAProtT5Extractor:
         emb_np = embedding.squeeze(0).cpu().numpy()
         return emb_np if emb_np.shape[0] > 0 else np.zeros((1, 1024), dtype=np.float32)
 
-# --- Predictor Model Architecture (Unchanged) ---
+
+# --- Predictor Model Head Architecture (Unchanged) ---
 class AntioxidantPredictor(nn.Module):
     def __init__(self, input_dim=1914, transformer_layers=3, transformer_heads=4, transformer_dropout=0.1):
         super(AntioxidantPredictor, self).__init__()
@@ -103,22 +102,34 @@ class AntioxidantPredictor(nn.Module):
     def get_temperature(self):
         return self.temperature.item()
 
-# --- Generator Model Architecture (Unchanged) ---
-class ProtT5Generator(nn.Module):
-    def __init__(self, vocab_size, embed_dim=512, num_layers=6, num_heads=8, dropout=0.1):
-        super(ProtT5Generator, self).__init__()
-        self.embed_tokens = nn.Embedding(vocab_size, embed_dim, padding_idx=token2id["<PAD>"])
-        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dropout=dropout, batch_first=True)
-        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+
+# --- ✅ NEW Generator Model Architecture ---
+class AdvancedProtT5Generator(nn.Module):
+    def __init__(self, base_model_id, 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)
+
+        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
         self.lm_head = nn.Linear(embed_dim, vocab_size)
+        
         self.vocab_size = vocab_size
         self.eos_token_id = token2id["<EOS>"]
         self.pad_token_id = token2id["<PAD>"]
+        print("  - Advanced Generator framework initialized.")
 
     def forward(self, input_ids):
-        embeddings = self.embed_tokens(input_ids)
-        encoder_output = self.encoder(embeddings)
-        logits = self.lm_head(encoder_output)
+        attention_mask = (input_ids != self.pad_token_id).int()
+        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
+        sequence_output = outputs.last_hidden_state
+        logits = self.lm_head(sequence_output)
         return logits
 
     def sample(self, batch_size, max_length=20, device="cpu", temperature=2.5, min_decoded_length=3):
@@ -132,6 +143,8 @@ class ProtT5Generator(nn.Module):
             probs = torch.softmax(next_logits, dim=-1)
             next_token = torch.multinomial(probs, num_samples=1)
             generated = torch.cat((generated, next_token), dim=1)
+            if (generated == self.eos_token_id).any(dim=1).all():
+                break
         return generated
 
     def decode(self, token_ids_batch):
@@ -140,8 +153,8 @@ class ProtT5Generator(nn.Module):
             seq = ""
             for token_id in ids_tensor.tolist()[1:]:
                 if token_id == self.eos_token_id: break
-                if token_id == self.pad_token_id: continue
-                seq += id2token.get(token_id, "")
+                if token_id == token2id["<PAD>"]: continue
+                seq += id2token.get(token_id, "?")
             sequences.append(seq)
         return sequences
 
@@ -151,34 +164,36 @@ try:
 except ImportError:
     raise gr.Error("Fatal Error: `feature_extract.py` not found. This file is required. Please upload it to your repository.")
 
-# --- Clustering Logic (Unchanged) ---
+# --- ✅ UPDATED Clustering Logic ---
 def cluster_sequences(generator, sequences, num_clusters, device):
     if not sequences or len(sequences) < num_clusters:
         return sequences[:num_clusters]
     with torch.no_grad():
         token_ids_list = []
-        max_len = max((len(seq) for seq in sequences), default=0) + 2
+        max_len = max(len(seq) for seq in sequences) + 2
         for seq in sequences:
-            ids = [token2id.get(aa, 0) for aa in seq] + [generator.eos_token_id]
-            ids = [np.random.randint(2, VOCAB_SIZE)] + ids
+            ids = [np.random.randint(2, VOCAB_SIZE)] + [token2id.get(aa, 0) for aa in seq] + [generator.eos_token_id]
             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)
-    representatives = []
+    reps = []
     for i in range(int(num_clusters)):
-        indices = np.where(kmeans.labels_ == i)[0]
-        if len(indices) == 0: continue
-        cluster_center = kmeans.cluster_centers_[i]
-        cluster_embeddings = seq_embeds_np[indices]
-        distances = np.linalg.norm(cluster_embeddings - cluster_center, axis=1)
-        representative_index = indices[np.argmin(distances)]
-        representatives.append(sequences[representative_index])
-    return representatives
+        idxs = np.where(kmeans.labels_ == i)[0]
+        if len(idxs) == 0: continue
+        center = kmeans.cluster_centers_[i]
+        distances = np.linalg.norm(seq_embeds_np[idxs] - center, axis=1)
+        rep_idx = idxs[np.argmin(distances)]
+        reps.append(sequences[rep_idx])
+    return reps
 
 # --------------------------------------------------------------------------
 # SECTION 2: GLOBAL MODEL AND DEPENDENCY LOADING
@@ -189,40 +204,43 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 try:
     # --- Define file paths ---
-    PREDICTOR_CHECKPOINT_PATH = "checkpoints/final_rl_model_logitp0.1_calibrated_FINETUNED_PROTT5.pth"
-    SCALER_PATH = "checkpoints/scaler_FINETUNED_PROTT5.pkl"
-    GENERATOR_CHECKPOINT_PATH = "generator_checkpoints_v3.6/final_generator_model.pth"
-    
-    # ✅ UPDATED: Define paths for LoRA-based loading
     PROTT5_BASE_MODEL_ID = "Rostlab/prot_t5_xl_uniref50"
-    LORA_ADAPTER_PATH = "./lora_finetuned_prott5" # Assumes LoRA files are in this directory
-
-    # --- Load Predictor Model (Head) ---
-    print(f"Loading Predictor from: {PREDICTOR_CHECKPOINT_PATH}")
+    
+    # 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")
+
+    # --- Load Validator System ---
+    print("--- Loading Validator System ---")
+    VALIDATOR_SCALER = joblib.load(SCALER_PATH)
+    VALIDATOR_EXTRACTOR = LoRAProtT5Extractor(
+        base_model_id=PROTT5_BASE_MODEL_ID,
+        lora_adapter_path=VALIDATOR_LORA_PATH
+    )
     PREDICTOR_MODEL = AntioxidantPredictor(input_dim=1914)
-    PREDICTOR_MODEL.load_state_dict(torch.load(PREDICTOR_CHECKPOINT_PATH, map_location=DEVICE))
+    PREDICTOR_MODEL.load_state_dict(torch.load(PREDICTOR_HEAD_CHECKPOINT_PATH, map_location=DEVICE))
     PREDICTOR_MODEL.to(DEVICE)
     PREDICTOR_MODEL.eval()
-    print(f"✅ Predictor model loaded (Temp: {PREDICTOR_MODEL.get_temperature():.4f}).")
-
-    # --- Load Scaler & LoRA Feature Extractor ---
-    print(f"Loading Scaler from: {SCALER_PATH}")
-    SCALER = joblib.load(SCALER_PATH)
-    print("Loading LoRA-enhanced ProtT5 Feature Extractor...")
-    # ✅ UPDATED: Instantiate the new LoRA extractor class
-    PROTT5_EXTRACTOR = LoRAProtT5Extractor(
+    print("✅ Validator System loaded successfully.")
+
+    # --- Load Generator System ---
+    print("\n--- Loading Generator System ---")
+    GENERATOR_MODEL = AdvancedProtT5Generator(
         base_model_id=PROTT5_BASE_MODEL_ID,
-        lora_adapter_path=LORA_ADAPTER_PATH
+        lora_adapter_path=GENERATOR_LORA_DIR,
+        vocab_size=VOCAB_SIZE
     )
-    print("✅ Scaler and Feature Extractor loaded.")
-
-    # --- Load Generator Model ---
-    print(f"Loading Generator from: {GENERATOR_CHECKPOINT_PATH}")
-    GENERATOR_MODEL = ProtT5Generator(vocab_size=VOCAB_SIZE)
-    GENERATOR_MODEL.load_state_dict(torch.load(GENERATOR_CHECKPOINT_PATH, map_location=DEVICE))
+    if not os.path.exists(GENERATOR_LM_HEAD_PATH):
+        raise FileNotFoundError(f"Generator's lm_head weights not found at: {GENERATOR_LM_HEAD_PATH}")
+    GENERATOR_MODEL.lm_head.load_state_dict(torch.load(GENERATOR_LM_HEAD_PATH, map_location=DEVICE))
     GENERATOR_MODEL.to(DEVICE)
     GENERATOR_MODEL.eval()
-    print("✅ Generator model loaded.")
+    print("✅ Generator System loaded successfully.")
     
     print("\n--- All models loaded! Gradio app is ready. ---\n")
 
@@ -231,7 +249,7 @@ except Exception as e:
     raise gr.Error(f"A required model or file could not be loaded. Please check your repository file structure and paths. Error details: {e}")
 
 # --------------------------------------------------------------------------
-# SECTION 3: WRAPPER FUNCTIONS FOR GRADIO UI (Unchanged logic)
+# SECTION 3: WRAPPER FUNCTIONS FOR GRADIO UI
 # --------------------------------------------------------------------------
 
 def predict_peptide_wrapper(sequence_str):
@@ -239,10 +257,9 @@ def predict_peptide_wrapper(sequence_str):
         return "0.0000", "Error: Please enter a valid peptide sequence using standard amino acids (ACDEFGHIKLMNPQRSTVWY)."
     
     try:
-        # This function call remains the same because the PROTT5_EXTRACTOR object,
-        # despite its new internal logic, provides the same interface.
-        features = extract_features(sequence_str.upper(), PROTT5_EXTRACTOR, L_fixed=29, d_model_pe=16)
-        scaled_features = SCALER.transform(features.reshape(1, -1))
+        # 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))
         
         with torch.no_grad():
             features_tensor = torch.tensor(scaled_features, dtype=torch.float32).to(DEVICE)
@@ -306,7 +323,7 @@ def generate_peptide_wrapper(num_to_generate, min_len, max_len, temperature, div
         time.sleep(1)
 
         if not validated_pool:
-            return pd.DataFrame([{"Sequence": "Could not generate any high-activity peptides (>0.9 prob) with the current settings. Try different parameters.", "Predicted Probability": "N/A"}])
+            return pd.DataFrame([{"Sequence": "Could not generate any high-activity peptides (>0.9 prob) with current settings.", "Predicted Probability": "N/A"}])
         
         high_quality_sequences = list(validated_pool.keys())
         final_diverse_seqs = cluster_sequences(GENERATOR_MODEL, high_quality_sequences, num_to_generate, DEVICE)