Update app.py

app.py

@@ -1,5 +1,4 @@
-# app.py - RLAnOxPeptide Gradio Web Application
-# This script integrates both the predictor and generator into a user-friendly web UI.
+# app.py - RLAnOxPeptide Gradio Web Application (Corrected Version)
 
 import os
 import torch
@@ -17,11 +16,10 @@ transformers.logging.set_verbosity_error()
 
 # --------------------------------------------------------------------------
 # SECTION 1: CORE CLASS AND FUNCTION DEFINITIONS
-# To make this app self-contained, we copy necessary class definitions here.
-# These should match the versions used during training.
+# These definitions are now synchronized with your provided, working scripts.
 # --------------------------------------------------------------------------
 
-# --- Vocabulary Definition (from both scripts) ---
+# --- Vocabulary Definition ---
 AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
 token2id = {aa: i + 2 for i, aa in enumerate(AMINO_ACIDS)}
 token2id["<PAD>"] = 0
@@ -29,16 +27,13 @@ token2id["<EOS>"] = 1
 id2token = {i: t for t, i in token2id.items()}
 VOCAB_SIZE = len(token2id)
 
-# --- Predictor Model Architecture (Corrected to match saved weights) ---
+# --- Predictor Model Architecture (VERSION THAT MATCHES YOUR .pth FILE) ---
 class AntioxidantPredictor(nn.Module):
     def __init__(self, input_dim, transformer_layers=3, transformer_heads=4, transformer_dropout=0.1):
         super(AntioxidantPredictor, self).__init__()
-        # 根据错误日志和您的训练脚本，我们知道输入维度是固定的
-        # 并且模型内部处理 ProtT5 和传统特征的分离
         self.t5_dim = 1024
         self.hand_crafted_dim = input_dim - self.t5_dim
         
-        # 定义 Transformer Encoder
         encoder_layer = nn.TransformerEncoderLayer(
             d_model=self.t5_dim, 
             nhead=transformer_heads, 
@@ -47,9 +42,6 @@ class AntioxidantPredictor(nn.Module):
         )
         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=transformer_layers)
         
-        # 定义 MLP
-        # 错误日志表明权重文件没有 fusion_fc 和 classifier，只有一个 mlp
-        # 我们根据 predictor_train.py 的原始结构来重建
         self.mlp = nn.Sequential(
             nn.Linear(input_dim, 512),
             nn.ReLU(),
@@ -62,19 +54,15 @@ class AntioxidantPredictor(nn.Module):
         self.temperature = nn.Parameter(torch.ones(1))
 
     def forward(self, fused_features):
-        # 这个前向传播逻辑与您的训练脚本 predictor_train.py 更为匹配
         prot_t5_features = fused_features[:, :self.t5_dim]
         hand_crafted_features = fused_features[:, self.t5_dim:]
         
-        # Transformer 只处理 ProtT5 特征
         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)
         
-        # 将处理后的 ProtT5 特征与传统特征拼接
         combined_features = torch.cat((transformer_output_pooled, hand_crafted_features), dim=1)
         
-        # 将最终拼接的特征送入 MLP
         logits = self.mlp(combined_features)
         
         return logits / self.temperature
@@ -87,7 +75,6 @@ class AntioxidantPredictor(nn.Module):
 
 # --- Generator Model Architecture (from generator.py) ---
 class ProtT5Generator(nn.Module):
-    # This class definition should be an exact copy from your project
     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>"])
@@ -97,7 +84,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)
@@ -112,12 +99,11 @@ class ProtT5Generator(nn.Module):
             next_logits = logits[:, -1, :] / temperature
             if generated.size(1) < min_decoded_length:
                 next_logits[:, self.eos_token_id] = -float("inf")
-            
             probs = torch.softmax(next_logits, dim=-1)
             next_token = torch.multinomial(probs, num_samples=1)
             generated = torch.cat((generated, next_token), dim=1)
-            
-            if (next_token == self.eos_token_id).all():
+            # 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
 
@@ -133,9 +119,7 @@ class ProtT5Generator(nn.Module):
             seqs.append(seq)
         return seqs
 
-# --- Feature Extraction Logic (from feature_extract.py) ---
-# Note: You need the actual ProtT5Model and extract_features here. 
-# Assuming they are in a file named `feature_extract.py` in the same directory.
+# --- Feature Extraction Logic (needs feature_extract.py) ---
 try:
     from feature_extract import ProtT5Model as FeatureProtT5Model, extract_features
 except ImportError:
@@ -147,11 +131,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 # Start token + EOS
+        max_len = max(len(seq) for seq in sequences) + 2 
         for seq in sequences:
-            # Recreate encoding to match how generator sees it (with start token)
             ids = [token2id.get(aa, 0) for aa in seq] + [generator.eos_token_id]
-            ids = [np.random.randint(2, VOCAB_SIZE)] + ids # Add a dummy start token
+            ids = [np.random.randint(2, VOCAB_SIZE)] + ids 
             ids += [token2id["<PAD>"]] * (max_len - len(ids))
             token_ids_list.append(ids)
 
@@ -175,55 +158,44 @@ def cluster_sequences(generator, sequences, num_clusters, device):
         representatives.append(sequences[representative_index])
     return representatives
 
-
 # --------------------------------------------------------------------------
 # SECTION 2: GLOBAL MODEL LOADING
-# Load all models and dependencies once when the app starts.
 # --------------------------------------------------------------------------
 print("Loading all models and dependencies. Please wait...")
-DEVICE = "cpu"  # Use CPU for compatibility with Hugging Face free tier
+DEVICE = "cpu"
 
 try:
-    # --- Define all required file paths here ---
-    # !! IMPORTANT: Ensure these are relative paths to the files in your Space !!
+    # --- Define file paths (!! CHECK THESE 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"
     PROTT5_BASE_MODEL_PATH = "prott5/model/"
     FINETUNED_PROTT5_FOR_FEATURES_PATH = "prott5/model/finetuned_prott5.bin"
 
-    # --- Load Predictor Components ---
+    # --- Load Predictor ---
     print("Loading Predictor Model...")
-    PREDICTOR_MODEL = AntioxidantPredictor(
-        input_dim=1914, transformer_layers=3, transformer_heads=4, transformer_dropout=0.1
-    )
+    PREDICTOR_MODEL = AntioxidantPredictor(input_dim=1914, transformer_layers=3, transformer_heads=4)
     PREDICTOR_MODEL.load_state_dict(torch.load(PREDICTOR_CHECKPOINT_PATH, map_location=DEVICE))
     PREDICTOR_MODEL.to(DEVICE)
     PREDICTOR_MODEL.eval()
-    print("✅ Predictor model loaded.")
+    print(f"✅ Predictor model loaded (Temp: {PREDICTOR_MODEL.get_temperature():.4f}).")
 
-    print("Loading Scaler...")
+    # --- Load Scaler & Feature Extractor ---
+    print("Loading Scaler and Feature Extractor...")
     SCALER = joblib.load(SCALER_PATH)
-    print("✅ Scaler loaded.")
-
-    print("Loading ProtT5 Feature Extractor...")
-    # This extractor must use the fine-tuned model for features, as per your training logic
     PROTT5_EXTRACTOR = FeatureProtT5Model(
         model_path=PROTT5_BASE_MODEL_PATH,
         finetuned_model_file=FINETUNED_PROTT5_FOR_FEATURES_PATH
     )
-    print("✅ ProtT5 Feature Extractor loaded.")
+    print("✅ Scaler and Feature Extractor loaded.")
 
-    # --- Load Generator Model ---
+    # --- Load Generator ---
     print("Loading Generator Model...")
-    GENERATOR_MODEL = ProtT5Generator(
-        vocab_size=VOCAB_SIZE, embed_dim=512, num_layers=6, num_heads=8, dropout=0.1
-    )
+    GENERATOR_MODEL = ProtT5Generator(vocab_size=VOCAB_SIZE, embed_dim=512, num_layers=6, num_heads=8)
     GENERATOR_MODEL.load_state_dict(torch.load(GENERATOR_CHECKPOINT_PATH, map_location=DEVICE))
     GENERATOR_MODEL.to(DEVICE)
     GENERATOR_MODEL.eval()
     print("✅ Generator model loaded.")
-
     print("\n--- All models loaded successfully! Gradio app is ready. ---\n")
 
 except Exception as e:
@@ -232,22 +204,17 @@ except Exception as e:
 
 # --------------------------------------------------------------------------
 # SECTION 3: WRAPPER FUNCTIONS FOR GRADIO
-# These functions connect the UI to our model's logic.
 # --------------------------------------------------------------------------
 
 def predict_peptide_wrapper(sequence_str):
-    """Takes a peptide sequence string and returns its predicted probability and class."""
     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 sequence with standard amino acids."
     
     try:
-        # 1. Extract features using the same logic as training/prediction scripts
-        features = extract_features(sequence_str, PROTT5_EXTRACTOR)
-        
-        # 2. Scale features
+        # 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))
         
-        # 3. Predict with the model
         with torch.no_grad():
             features_tensor = torch.tensor(scaled_features, dtype=torch.float32).to(DEVICE)
             logits = PREDICTOR_MODEL(features_tensor)
@@ -261,47 +228,46 @@ 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)):
-    """Generates, validates, and clusters sequences."""
     num_to_generate = int(num_to_generate)
     min_len = int(min_len)
     max_len = int(max_len)
     
     try:
-        # STEP 1: Generate an initial pool of unique sequences
+        # Step 1: Generate a pool of unique sequences
         target_pool_size = int(num_to_generate * diversity_factor)
         unique_seqs = set()
-        progress(0, desc="Generating initial peptide pool...")
         
-        max_attempts = 10
-        attempts = 0
-        while len(unique_seqs) < target_pool_size and attempts < max_attempts:
-            batch_size = (target_pool_size - len(unique_seqs)) * 2 # Generate extra to account for duplicates/short ones
-            with torch.no_grad():
-                generated_tokens = GENERATOR_MODEL.sample(
-                    batch_size=batch_size, 
-                    max_length=max_len, 
-                    device=DEVICE,
-                    temperature=temperature, 
-                    min_decoded_length=min_len
-                )
-            decoded = GENERATOR_MODEL.decode(generated_tokens.cpu())
-            for seq in decoded:
-                if min_len <= len(seq) <= max_len:
-                    unique_seqs.add(seq)
-            attempts += 1
-            progress(len(unique_seqs) / target_pool_size, desc=f"Generated {len(unique_seqs)} unique sequences...")
+        # 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))
+                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
+                    )
+                decoded = GENERATOR_MODEL.decode(generated_tokens.cpu())
+                
+                newly_added = 0
+                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)
 
         candidate_seqs = list(unique_seqs)
-        if not candidate_seqs:
-            return pd.DataFrame({"Sequence": ["Failed to generate valid sequences."], "Predicted Probability": ["N/A"]})
 
-        # STEP 2: Validate the generated sequences
+        # Step 2: Validate the generated sequences
         validated_pool = {}
         for seq in tqdm(candidate_seqs, desc="Validating generated sequences"):
             prob_str, _ = predict_peptide_wrapper(seq)
             try:
                 prob = float(prob_str)
-                if prob > 0.90: # Filter for high-quality peptides as in generator.py
+                if prob > 0.90:  # Filter for high-quality peptides
                     validated_pool[seq] = prob
             except (ValueError, TypeError):
                 continue
@@ -311,11 +277,11 @@ def generate_peptide_wrapper(num_to_generate, min_len, max_len, temperature, div
         
         high_quality_sequences = list(validated_pool.keys())
         
-        # STEP 3: Cluster to ensure diversity
+        # Step 3: Cluster to ensure diversity
         progress(1.0, desc="Clustering for diversity...")
         final_diverse_seqs = cluster_sequences(GENERATOR_MODEL, high_quality_sequences, num_to_generate, DEVICE)
         
-        # STEP 4: Format final results
+        # Step 4: Format final results
         final_results = [(seq, f"{validated_pool[seq]:.4f}") for seq in final_diverse_seqs]
         final_results.sort(key=lambda x: float(x[1]), reverse=True)
         
@@ -325,10 +291,8 @@ def generate_peptide_wrapper(num_to_generate, min_len, max_len, temperature, div
         print(f"Generation error: {e}")
         return pd.DataFrame({"Sequence": [f"An error occurred during generation: {e}"], "Predicted Probability": ["N/A"]})
 
-
 # --------------------------------------------------------------------------
-# SECTION 4: GRADIO UI CONSTRUCTION
-# Building the web interface. All text is in English.
+# SECTION 4: GRADIO UI CONSTRUCTION (ALL ENGLISH)
 # --------------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft(), title="RLAnOxPeptide") as demo:
     gr.Markdown("# RLAnOxPeptide: Intelligent Peptide Design and Prediction Platform")
@@ -350,11 +314,8 @@ with gr.Blocks(theme=gr.themes.Soft(), title="RLAnOxPeptide") as demo:
                 outputs=[probability_output, class_output]
             )
             gr.Examples(
-                examples=[["WHYHDYKY"], ["YPGG"], ["LVLHEHGGN"], ["INVALIDSEQUENCE"]],
-                inputs=peptide_input,
-                outputs=[probability_output, class_output],
-                fn=predict_peptide_wrapper,
-                cache_examples=False,
+                examples=[["WHYHDYKY"], ["YPGG"], ["LVLHEHGGN"]],
+                inputs=peptide_input
             )
 
         with gr.TabItem("Novel Sequence Generator"):
@@ -378,4 +339,4 @@ with gr.Blocks(theme=gr.themes.Soft(), title="RLAnOxPeptide") as demo:
             )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()