Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py - RLAnOxPeptide Gradio Web Application (Corrected Version)
+# app.py - RLAnOxPeptide Gradio Web Application (FINAL CORRECTED VERSION)
 
 import os
 import torch
@@ -27,53 +27,66 @@ token2id["<EOS>"] = 1
 id2token = {i: t for t, i in token2id.items()}
 VOCAB_SIZE = len(token2id)
 
-# --- Predictor Model Architecture (VERSION THAT MATCHES YOUR .pth FILE) ---
+# --- Predictor Model Architecture (Copied from your LATEST 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__()
-        self.t5_dim = 1024
-        self.hand_crafted_dim = input_dim - self.t5_dim
-        
+        self.prott5_dim = 1024
+        self.handcrafted_dim = input_dim - self.prott5_dim
+        self.seq_len = 16
+        self.prott5_feature_dim = 64  # 16 * 64 = 1024
+
         encoder_layer = nn.TransformerEncoderLayer(
-            d_model=self.t5_dim, 
-            nhead=transformer_heads, 
-            dropout=transformer_dropout, 
+            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)
-        
-        self.mlp = nn.Sequential(
-            nn.Linear(input_dim, 512),
+
+        fused_dim = self.prott5_feature_dim + self.handcrafted_dim
+        self.fusion_fc = nn.Sequential(
+            nn.Linear(fused_dim, 1024),
             nn.ReLU(),
-            nn.Dropout(0.5),
+            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.5),
+            nn.Dropout(0.3),
             nn.Linear(256, 1)
         )
-        self.temperature = nn.Parameter(torch.ones(1))
+        self.temperature = nn.Parameter(torch.ones(1), requires_grad=False)
 
-    def forward(self, fused_features):
-        prot_t5_features = fused_features[:, :self.t5_dim]
-        hand_crafted_features = fused_features[:, self.t5_dim:]
+    def forward(self, x, *args):
+        batch_size = x.size(0)
+        prot_t5_features = x[:, :self.prott5_dim]
+        handcrafted_features = x[:, self.prott5_dim:]
         
-        prot_t5_features_unsqueezed = prot_t5_features.unsqueeze(1)
-        transformer_output = self.transformer_encoder(prot_t5_features_unsqueezed)
-        transformer_output_pooled = transformer_output.mean(dim=1)
+        prot_t5_seq = prot_t5_features.view(batch_size, self.seq_len, self.prott5_feature_dim)
+        encoded_seq = self.transformer_encoder(prot_t5_seq)
+        refined_prott5 = encoded_seq.mean(dim=1)
         
-        combined_features = torch.cat((transformer_output_pooled, hand_crafted_features), dim=1)
+        fused_features = torch.cat([refined_prott5, handcrafted_features], dim=1)
+        fused_features = self.fusion_fc(fused_features)
         
-        logits = self.mlp(combined_features)
+        logits = self.classifier(fused_features)
         
-        return logits / self.temperature
-
-    def get_temperature(self):
-        return self.temperature.item()
+        logits_scaled = logits / self.temperature
+        
+        return logits_scaled
 
     def set_temperature(self, temp_value, device):
         self.temperature = nn.Parameter(torch.tensor([temp_value], device=device), requires_grad=False)
 
-# --- Generator Model Architecture (from generator.py) ---
+    def get_temperature(self):
+        return self.temperature.item()
+
+# --- Generator Model Architecture (Copied 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__()
@@ -102,7 +115,6 @@ 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)
-            # Early stop if all sequences in batch have generated an EOS token
             if (generated == self.eos_token_id).any(dim=1).all():
                 break
         return generated
@@ -111,19 +123,18 @@ class ProtT5Generator(nn.Module):
         seqs = []
         for ids_tensor in token_ids_batch:
             seq = ""
-            # Start from index 1 to skip the initial random start token
-            for token_id in ids_tensor.tolist()[1:]:
+            for token_id in ids_tensor.tolist()[1:]: # Skip start token
                 if token_id == self.eos_token_id: break
                 if token_id == self.pad_token_id: continue
                 seq += id2token.get(token_id, "?")
             seqs.append(seq)
         return seqs
 
-# --- Feature Extraction Logic (needs feature_extract.py) ---
+# --- Feature Extraction (needs feature_extract.py) ---
 try:
     from feature_extract import ProtT5Model as FeatureProtT5Model, extract_features
 except ImportError:
-    raise gr.Error("Failed to import feature_extract.py. Please ensure the file is in the same directory as app.py.")
+    raise gr.Error("Failed to import feature_extract.py. Ensure it is in the same directory.")
 
 # --- Clustering Logic (from generator.py) ---
 def cluster_sequences(generator, sequences, num_clusters, device):
@@ -131,10 +142,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)
 
@@ -161,7 +172,7 @@ def cluster_sequences(generator, sequences, num_clusters, device):
 # --------------------------------------------------------------------------
 # SECTION 2: GLOBAL MODEL LOADING
 # --------------------------------------------------------------------------
-print("Loading all models and dependencies. Please wait...")
+print("Loading all models and dependencies...")
 DEVICE = "cpu"
 
 try:
@@ -174,7 +185,9 @@ try:
 
     # --- Load Predictor ---
     print("Loading Predictor Model...")
-    PREDICTOR_MODEL = AntioxidantPredictor(input_dim=1914, transformer_layers=3, transformer_heads=4)
+    PREDICTOR_MODEL = AntioxidantPredictor(
+        input_dim=1914, transformer_layers=3, transformer_heads=4, transformer_dropout=0.1
+    )
     PREDICTOR_MODEL.load_state_dict(torch.load(PREDICTOR_CHECKPOINT_PATH, map_location=DEVICE))
     PREDICTOR_MODEL.to(DEVICE)
     PREDICTOR_MODEL.eval()
@@ -191,7 +204,9 @@ try:
 
     # --- Load Generator ---
     print("Loading Generator Model...")
-    GENERATOR_MODEL = ProtT5Generator(vocab_size=VOCAB_SIZE, embed_dim=512, num_layers=6, num_heads=8)
+    GENERATOR_MODEL = ProtT5Generator(
+        vocab_size=VOCAB_SIZE, embed_dim=512, num_layers=6, num_heads=8, dropout=0.1
+    )
     GENERATOR_MODEL.load_state_dict(torch.load(GENERATOR_CHECKPOINT_PATH, map_location=DEVICE))
     GENERATOR_MODEL.to(DEVICE)
     GENERATOR_MODEL.eval()
@@ -211,7 +226,7 @@ def predict_peptide_wrapper(sequence_str):
         return "0.0000", "Error: Please enter a valid sequence with standard amino acids."
     
     try:
-        # These L_fixed and d_model_pe values are from your predictor.py args
+        # Use feature extraction params from your working predictor.py
         features = extract_features(sequence_str, PROTT5_EXTRACTOR, L_fixed=29, d_model_pe=16)
         scaled_features = SCALER.transform(features.reshape(1, -1))
         
@@ -237,27 +252,22 @@ def generate_peptide_wrapper(num_to_generate, min_len, max_len, temperature, div
         target_pool_size = int(num_to_generate * diversity_factor)
         unique_seqs = set()
         
-        # A simple generation loop based on generator.py logic
         with tqdm(total=target_pool_size, desc="Generating candidate sequences") as pbar:
             while len(unique_seqs) < target_pool_size:
-                batch_size = (target_pool_size - len(unique_seqs))
+                # Generate a surplus to account for filtering
+                batch_size = max(1, (target_pool_size - len(unique_seqs)) * 2)
                 with torch.no_grad():
                     generated_tokens = GENERATOR_MODEL.sample(
-                        batch_size=max(1, batch_size), 
-                        max_length=max_len, 
-                        device=DEVICE,
-                        temperature=temperature, 
-                        min_decoded_length=min_len
+                        batch_size=batch_size, max_length=max_len, device=DEVICE,
+                        temperature=temperature, min_decoded_length=min_len
                     )
                 decoded = GENERATOR_MODEL.decode(generated_tokens.cpu())
                 
-                newly_added = 0
+                initial_count = len(unique_seqs)
                 for seq in decoded:
                     if min_len <= len(seq) <= max_len:
-                        if seq not in unique_seqs:
-                           unique_seqs.add(seq)
-                           newly_added +=1
-                pbar.update(newly_added)
+                        unique_seqs.add(seq)
+                pbar.update(len(unique_seqs) - initial_count)
 
         candidate_seqs = list(unique_seqs)
 
@@ -267,13 +277,13 @@ def generate_peptide_wrapper(num_to_generate, min_len, max_len, temperature, div
             prob_str, _ = predict_peptide_wrapper(seq)
             try:
                 prob = float(prob_str)
-                if prob > 0.90:  # Filter for high-quality peptides
+                if prob > 0.90:
                     validated_pool[seq] = prob
             except (ValueError, TypeError):
                 continue
 
         if not validated_pool:
-            return pd.DataFrame({"Sequence": ["No high-activity peptides (>0.9 prob) were generated."], "Predicted Probability": ["N/A"]})
+            return pd.DataFrame([{"Sequence": "No high-activity peptides (>0.9 prob) were generated.", "Predicted Probability": "N/A"}])
         
         high_quality_sequences = list(validated_pool.keys())
         
@@ -289,10 +299,10 @@ def generate_peptide_wrapper(num_to_generate, min_len, max_len, temperature, div
 
     except Exception as e:
         print(f"Generation error: {e}")
-        return pd.DataFrame({"Sequence": [f"An error occurred during generation: {e}"], "Predicted Probability": ["N/A"]})
+        return pd.DataFrame([{"Sequence": f"An error occurred: {e}", "Predicted Probability": "N/A"}])
 
 # --------------------------------------------------------------------------
-# SECTION 4: GRADIO UI CONSTRUCTION (ALL ENGLISH)
+# SECTION 4: GRADIO UI CONSTRUCTION
 # --------------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft(), title="RLAnOxPeptide") as demo:
     gr.Markdown("# RLAnOxPeptide: Intelligent Peptide Design and Prediction Platform")