Update app.py

app.py

@@ -1,4 +1,5 @@
-# app.py - RLAnOxPeptide Gradio Web Application (FINAL CORRECTED VERSION - Synced with local scripts)
+
+# app.py - RLAnOxPeptide Gradio Web Application (FINAL CORRECTED VERSION - Robust Loading)
 
 import os
 import torch
@@ -10,17 +11,15 @@ 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()
 
 # --------------------------------------------------------------------------
 # SECTION 1: CORE CLASS AND FUNCTION DEFINITIONS
-# These definitions are now synchronized with your provided, working scripts.
 # --------------------------------------------------------------------------
 
-# --- Vocabulary Definition (from generator.py) ---
+# --- Vocabulary Definition ---
 AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
 token2id = {aa: i + 2 for i, aa in enumerate(AMINO_ACIDS)}
 token2id["<PAD>"] = 0
@@ -28,62 +27,69 @@ token2id["<EOS>"] = 1
 id2token = {i: t for t, i in token2id.items()}
 VOCAB_SIZE = len(token2id)
 
-# --- Predictor Model Architecture (Copied VERBATIM from your antioxidant_predictor_5.py) ---
+# --- ROBUST FeatureProtT5Model Class for Feature Extraction ---
+class FeatureProtT5Model:
+    def __init__(self, model_dir_path, finetuned_weights_path=None):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Initializing ProtT5 from base directory: {model_dir_path}")
+        
+        # Step 1: Load the base model architecture and tokenizer from the directory.
+        # This step requires the original pytorch_model.bin to be in the model_dir_path.
+        self.tokenizer = transformers.T5Tokenizer.from_pretrained(model_dir_path, do_lower_case=False)
+        self.model = transformers.T5EncoderModel.from_pretrained(model_dir_path)
+        
+        # Step 2: If a separate fine-tuned weights file is provided, load it.
+        if finetuned_weights_path and os.path.exists(finetuned_weights_path):
+            print(f"Loading and applying fine-tuned weights from: {finetuned_weights_path}")
+            # Load the state_dict from your specific fine-tuned file
+            state_dict = torch.load(finetuned_weights_path, map_location=self.device)
+            # Use strict=False because the fine-tuned model may only contain encoder weights
+            self.model.load_state_dict(state_dict, strict=False)
+            print("Successfully applied fine-tuned weights to the model.")
+        else:
+            print("Warning: Fine-tuned weights file not provided or not found. Using the base ProtT5 model weights.")
+            
+        self.model.to(self.device)
+        self.model.eval()
+
+    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, max_length=1022)
+        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 = embedding.squeeze(0).cpu().numpy()
+        return emb if emb.shape[0] > 0 else np.zeros((1, 1024), dtype=np.float32)
+
+# --- Predictor Model Architecture ---
 class AntioxidantPredictor(nn.Module):
     def __init__(self, input_dim, transformer_layers=3, transformer_heads=4, transformer_dropout=0.1):
         super(AntioxidantPredictor, self).__init__()
         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.prott5_feature_dim,
-            nhead=transformer_heads,
-            dropout=transformer_dropout,
-            batch_first=True
-        )
+        self.prott5_feature_dim = 64
+        encoder_layer = nn.TransformerEncoderLayer(d_model=self.prott5_feature_dim, nhead=transformer_heads, dropout=transformer_dropout, batch_first=.T.rue)
         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.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, *args):
         batch_size = x.size(0)
         prot_t5_features = x[:, :self.prott5_dim]
         handcrafted_features = x[:, self.prott5_dim:]
-        
         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)
-        
         fused_features = torch.cat([refined_prott5, handcrafted_features], dim=1)
         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):
-        self.temperature = nn.Parameter(torch.tensor([temp_value], device=device), requires_grad=False)
-
-    def get_temperature(self):
-        return self.temperature.item()
+        return logits / self.temperature
+    def set_temperature(self, temp_value, device): self.temperature = nn.Parameter(torch.tensor([temp_value], device=device), requires_grad=False)
+    def get_temperature(self): return self.temperature.item()
 
 # --- Generator Model Architecture (Copied VERBATIM from your generator.py) ---
 class ProtT5Generator(nn.Module):