Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py - RLAnOxPeptide Gradio Web Application (FINAL CORRECTED VERSION)
+# app.py - RLAnOxPeptide Gradio Web Application (FINAL CORRECTED VERSION - Synced with local scripts)
 
 import os
 import torch
@@ -10,6 +10,7 @@ import gradio as gr
 from sklearn.cluster import KMeans
 from tqdm import tqdm
 import transformers
+import argparse # We won't use argparse but might need it for compatibility if any function expects it
 
 # Suppress verbose logging from transformers
 transformers.logging.set_verbosity_error()
@@ -19,7 +20,7 @@ transformers.logging.set_verbosity_error()
 # These definitions are now synchronized with your provided, working scripts.
 # --------------------------------------------------------------------------
 
-# --- Vocabulary Definition ---
+# --- Vocabulary Definition (from generator.py) ---
 AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
 token2id = {aa: i + 2 for i, aa in enumerate(AMINO_ACIDS)}
 token2id["<PAD>"] = 0
@@ -27,7 +28,7 @@ token2id["<EOS>"] = 1
 id2token = {i: t for t, i in token2id.items()}
 VOCAB_SIZE = len(token2id)
 
-# --- Predictor Model Architecture (Copied from your LATEST antioxidant_predictor_5.py) ---
+# --- Predictor Model Architecture (Copied VERBATIM from your antioxidant_predictor_5.py) ---
 class AntioxidantPredictor(nn.Module):
     def __init__(self, input_dim, transformer_layers=3, transformer_heads=4, transformer_dropout=0.1):
         super(AntioxidantPredictor, self).__init__()
@@ -75,9 +76,7 @@ class AntioxidantPredictor(nn.Module):
         fused_features = self.fusion_fc(fused_features)
         
         logits = self.classifier(fused_features)
-        
         logits_scaled = logits / self.temperature
-        
         return logits_scaled
 
     def set_temperature(self, temp_value, device):
@@ -86,7 +85,7 @@ class AntioxidantPredictor(nn.Module):
     def get_temperature(self):
         return self.temperature.item()
 
-# --- Generator Model Architecture (Copied from your generator.py) ---
+# --- Generator Model Architecture (Copied VERBATIM from your generator.py) ---
 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__()
@@ -97,7 +96,7 @@ class ProtT5Generator(nn.Module):
         self.vocab_size = vocab_size
         self.eos_token_id = token2id["<EOS>"]
         self.pad_token_id = token2id["<PAD>"]
-
+    
     def forward(self, input_ids):
         embeddings = self.embed_tokens(input_ids)
         encoder_output = self.encoder(embeddings)
@@ -130,7 +129,7 @@ class ProtT5Generator(nn.Module):
             seqs.append(seq)
         return seqs
 
-# --- Feature Extraction (needs feature_extract.py) ---
+# --- Feature Extraction (needs feature_extract.py in the same directory) ---
 try:
     from feature_extract import ProtT5Model as FeatureProtT5Model, extract_features
 except ImportError:
@@ -142,10 +141,10 @@ def cluster_sequences(generator, sequences, num_clusters, device):
         return sequences[:num_clusters]
     with torch.no_grad():
         token_ids_list = []
-        max_len = max(len(seq) for seq in sequences) + 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)] + ids 
             ids += [token2id["<PAD>"]] * (max_len - len(ids))
             token_ids_list.append(ids)
 
@@ -181,13 +180,16 @@ try:
     SCALER_PATH = "checkpoints/scaler_FINETUNED_PROTT5.pkl"
     GENERATOR_CHECKPOINT_PATH = "generator_checkpoints_v3.6/final_generator_model.pth"
     PROTT5_BASE_MODEL_PATH = "prott5/model/"
+    # This path is now used by the FeatureProtT5Model to load the fine-tuned weights
     FINETUNED_PROTT5_FOR_FEATURES_PATH = "prott5/model/finetuned_prott5.bin"
 
     # --- Load Predictor ---
     print("Loading Predictor Model...")
+    # Initialize the correct class
     PREDICTOR_MODEL = AntioxidantPredictor(
         input_dim=1914, transformer_layers=3, transformer_heads=4, transformer_dropout=0.1
     )
+    # Load the state dict that matches this class
     PREDICTOR_MODEL.load_state_dict(torch.load(PREDICTOR_CHECKPOINT_PATH, map_location=DEVICE))
     PREDICTOR_MODEL.to(DEVICE)
     PREDICTOR_MODEL.eval()
@@ -226,7 +228,7 @@ def predict_peptide_wrapper(sequence_str):
         return "0.0000", "Error: Please enter a valid sequence with standard amino acids."
     
     try:
-        # Use feature extraction params from your working predictor.py
+        # These L_fixed and d_model_pe values are from your predictor.py args
         features = extract_features(sequence_str, PROTT5_EXTRACTOR, L_fixed=29, d_model_pe=16)
         scaled_features = SCALER.transform(features.reshape(1, -1))
         
@@ -243,6 +245,7 @@ def predict_peptide_wrapper(sequence_str):
         return "N/A", f"An error occurred during processing: {e}"
 
 def generate_peptide_wrapper(num_to_generate, min_len, max_len, temperature, diversity_factor, progress=gr.Progress(track_tqdm=True)):
+    # This logic is a direct adaptation of your generator.py main function
     num_to_generate = int(num_to_generate)
     min_len = int(min_len)
     max_len = int(max_len)
@@ -254,8 +257,7 @@ def generate_peptide_wrapper(num_to_generate, min_len, max_len, temperature, div
         
         with tqdm(total=target_pool_size, desc="Generating candidate sequences") as pbar:
             while len(unique_seqs) < target_pool_size:
-                # Generate a surplus to account for filtering
-                batch_size = max(1, (target_pool_size - len(unique_seqs)) * 2)
+                batch_size = max(1, (target_pool_size - len(unique_seqs)))
                 with torch.no_grad():
                     generated_tokens = GENERATOR_MODEL.sample(
                         batch_size=batch_size, max_length=max_len, device=DEVICE,