app.py

# app.py - RLAnOxPeptide Gradio Web Application
# This script integrates both the predictor and generator into a user-friendly web UI.

import os
import torch
import torch.nn as nn
import pandas as pd
import joblib
import numpy as np
import gradio as gr
from sklearn.cluster import KMeans
from tqdm import tqdm
import transformers

# Suppress verbose logging from transformers
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.
# --------------------------------------------------------------------------

# --- Vocabulary Definition (from both scripts) ---
AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
token2id = {aa: i + 2 for i, aa in enumerate(AMINO_ACIDS)}
token2id["<PAD>"] = 0
token2id["<EOS>"] = 1
id2token = {i: t for t, i in token2id.items()}
VOCAB_SIZE = len(token2id)

# --- Predictor Model Architecture (from antioxidant_predictor_5.py) ---
class AntioxidantPredictor(nn.Module):
    # This class definition should be an exact copy from your project
    def __init__(self, input_dim, transformer_layers=3, transformer_heads=4, transformer_dropout=0.1):
        super(AntioxidantPredictor, self).__init__()
        self.input_dim = input_dim
        self.t5_dim = 1024
        self.hand_crafted_dim = self.input_dim - self.t5_dim
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=self.t5_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(self.input_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, 1)
        )
        self.temperature = nn.Parameter(torch.ones(1))

    def forward(self, fused_features):
        tr_features = fused_features[:, :self.t5_dim]
        hand_features = fused_features[:, self.t5_dim:]
        tr_features_unsqueezed = tr_features.unsqueeze(1)
        transformer_output = self.transformer_encoder(tr_features_unsqueezed)
        transformer_output_pooled = transformer_output.mean(dim=1)
        combined_features = torch.cat((transformer_output_pooled, hand_features), dim=1)
        logits = self.mlp(combined_features)
        return logits / self.temperature

    def get_temperature(self):
        return self.temperature.item()

# --- 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>"])
        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)
        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>"]
    
    def forward(self, input_ids):
        embeddings = self.embed_tokens(input_ids)
        encoder_output = self.encoder(embeddings)
        logits = self.lm_head(encoder_output)
        return logits
    
    def sample(self, batch_size, max_length=20, device="cpu", temperature=2.5, min_decoded_length=3):
        start_token = torch.randint(2, self.vocab_size, (batch_size, 1), device=device)
        generated = start_token
        for _ in range(max_length - 1):
            logits = self.forward(generated)
            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():
                break
        return generated

    def decode(self, token_ids_batch):
        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:]:
                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 (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.
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.")

# --- Clustering Logic (from generator.py) ---
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) + 2 # Start token + EOS
        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 += [token2id["<PAD>"]] * (max_len - len(ids))
            token_ids_list.append(ids)

        input_ids = torch.tensor(token_ids_list, dtype=torch.long, device=device)
        embeddings = generator.embed_tokens(input_ids)
        mask = (input_ids != token2id["<PAD>"]).unsqueeze(-1).float()
        embeddings = embeddings * mask
        lengths = mask.sum(dim=1)
        seq_embeds = embeddings.sum(dim=1) / (lengths + 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 = []
    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


# --------------------------------------------------------------------------
# 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

try:
    # --- Define all required file paths here ---
    # !! IMPORTANT: Ensure these are relative paths to the files in your Space !!
    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 ---
    print("Loading Predictor Model...")
    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()
    print("✅ Predictor model loaded.")

    print("Loading Scaler...")
    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.")

    # --- Load Generator Model ---
    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.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:
    print(f"💥 FATAL ERROR: Failed to load a model or dependency file: {e}")
    raise gr.Error(f"Model or dependency loading failed! Check file paths and integrity. Error: {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
        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)
            probability = torch.sigmoid(logits).squeeze().item()
        
        classification = "Antioxidant" if probability >= 0.5 else "Non-Antioxidant"
        return f"{probability:.4f}", classification

    except Exception as e:
        print(f"Prediction error for sequence '{sequence_str}': {e}")
        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
        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...")

        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
        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
                    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"]})
        
        high_quality_sequences = list(validated_pool.keys())
        
        # 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
        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)
        
        return pd.DataFrame(final_results, columns=["Sequence", "Predicted Probability"])

    except Exception as e:
        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.
# --------------------------------------------------------------------------
with gr.Blocks(theme=gr.themes.Soft(), title="RLAnOxPeptide") as demo:
    gr.Markdown("# RLAnOxPeptide: Intelligent Peptide Design and Prediction Platform")
    gr.Markdown("An integrated framework combining reinforcement learning and a Transformer model for the efficient prediction and innovative design of antioxidant peptides.")

    with gr.Tabs():
        with gr.TabItem("Peptide Activity Predictor"):
            gr.Markdown("### Enter an amino acid sequence to predict its antioxidant activity.")
            with gr.Row():
                peptide_input = gr.Textbox(label="Peptide Sequence", placeholder="e.g., WHYHDYKY", scale=3)
                predict_button = gr.Button("Predict", variant="primary", scale=1)
            with gr.Row():
                probability_output = gr.Textbox(label="Predicted Probability")
                class_output = gr.Textbox(label="Predicted Class")
            
            predict_button.click(
                fn=predict_peptide_wrapper,
                inputs=peptide_input,
                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,
            )

        with gr.TabItem("Novel Sequence Generator"):
            gr.Markdown("### Set parameters to generate novel, high-activity antioxidant peptides.")
            with gr.Column():
                with gr.Row():
                    num_input = gr.Slider(minimum=1, maximum=50, value=10, step=1, label="Number of Final Peptides to Generate")
                    min_len_input = gr.Slider(minimum=2, maximum=10, value=3, step=1, label="Minimum Length")
                    max_len_input = gr.Slider(minimum=10, maximum=20, value=20, step=1, label="Maximum Length")
                with gr.Row():
                    temp_input = gr.Slider(minimum=0.5, maximum=3.0, value=2.5, step=0.1, label="Temperature (Higher = More random)")
                    diversity_input = gr.Slider(minimum=1.0, maximum=3.0, value=1.2, step=0.1, label="Diversity Factor (Higher = Larger initial pool for clustering)")
            
            generate_button = gr.Button("Generate Peptides", variant="primary")
            results_output = gr.DataFrame(headers=["Sequence", "Predicted Probability"], label="Generated & Validated Peptides", wrap=True)

            generate_button.click(
                fn=generate_peptide_wrapper,
                inputs=[num_input, min_len_input, max_len_input, temp_input, diversity_input],
                outputs=results_output
            )

if __name__ == "__main__":
    demo.launch()