Update app.py

app.py

@@ -1,41 +1,116 @@
+import os
+import urllib.request
 import gradio as gr
 import torch
-from huggingface_hub import hf_hub_download
-# import prody # (Uncomment when you add your parsing logic)
+import prody
+import numpy as np
 
-# 1. Download the weights from your Model repository
-# Make sure "WSobo/Struct2Seq-GNN" matches exactly what you named your model repo
-model_path = hf_hub_download(repo_id="WSobo/Struct2Seq-GNN", filename="best_model.pt")
+# Import your model class (Make sure model_utils.py is uploaded to your Space!)
+from model_utils import Struct2SeqGNN 
 
-# 2. Load the model (Ensure your model class is defined in this file or imported)
-# model = MyGNNModel()
-# model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
-# model.eval()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+# ---------------------------------------------------------
+# 1. DOWNLOAD & LOAD MODEL WEIGHTS
+# ---------------------------------------------------------
+raw_github_url = "https://raw.githubusercontent.com/WSobo/Struct2Seq-GNN/main/pretrained_models/v2.0/best_model.pt"
+model_path = "best_model.pt"
+
+# Download weights if they aren't already cached in the Space
+if not os.path.exists(model_path):
+    print("Downloading model weights from GitHub...")
+    urllib.request.urlretrieve(raw_github_url, model_path)
+
+# Instantiate the model matching your v2.0 training parameters
+model = Struct2SeqGNN(
+    node_features=6, 
+    ligand_features=6, 
+    hidden_dim=256, 
+    num_classes=21, 
+    num_layers=6, 
+    dropout=0.0
+).to(device)
+
+# Load the weights with DDP prefix handling
+state_dict = torch.load(model_path, map_location=device)
+if list(state_dict.keys())[0].startswith('module.'):
+    state_dict = {k[7:]: v for k, v in state_dict.items()}
+model.load_state_dict(state_dict)
+model.eval()
+
+# Standard Amino Acid alphabet (Update this if your model uses a different index order!)
+AA_ALPHABET = "ACDEFGHIKLMNPQRSTVWYX"
+
+
+# ---------------------------------------------------------
+# 2. INFERENCE PIPELINE
+# ---------------------------------------------------------
 def predict_sequence(pdb_file):
-    """
-    This function is triggered every time a user hits 'Submit'.
-    """
     if pdb_file is None:
         return "Please upload a .pdb file."
     
-    # -- YOUR INFERENCE PIPELINE GOES HERE --
-    # 1. Parse the uploaded pdb_file.name using ProDy
-    # 2. Convert coordinates to tensors
-    # 3. Run the forward pass: logits = model(coords)
-    # 4. Decode logits to an amino acid string
-    
-    # Dummy return for testing the UI
-    return "MKTLLILAVIM... (Your Struct2Seq-GNN prediction will appear here!)"
+    try:
+        # Step 2.1: Silence ProDy logs to keep your server console clean
+        prody.confProDy(verbosity='none')
+        
+        # Step 2.2: Parse the uploaded PDB
+        pdb = prody.parsePDB(pdb_file.name)
+        if pdb is None:
+            return "Error: Could not parse the PDB file."
+            
+        # Step 2.3: Extract backbone coordinates
+        # (Grabbing C-alphas to get the sequence length and main coordinates)
+        calphas = pdb.select('calpha')
+        if calphas is None:
+            return "Error: No alpha carbons found in the PDB."
+        
+        num_residues = len(calphas)
+        
+        # Step 2.4: Convert coordinates to numpy/tensors
+        # coords = calphas.getCoords()  # Shape: [num_residues, 3]
+        
+        # =====================================================================
+        # Step 2.5: YOUR GRAPH CONSTRUCTION GOES HERE
+        # Copy the exact logic you use in `utils.dataset.Struct2SeqDataset` 
+        # to turn these coordinates into your graph components (x, edge_index).
+        # =====================================================================
+        
+        # NOTE: Delete this placeholder dummy block once your logic is in!
+        dummy_logits = torch.randn((num_residues, 21)).to(device)
+        
+        # Step 2.6: Run the forward pass
+        with torch.no_grad():
+            # logits = model(x, edge_index, ...) 
+            logits = dummy_logits # Placeholder
+            
+        # Step 2.7: Decode logits to an amino acid string
+        # Argmax gets the index of the highest probability AA for each residue
+        predicted_indices = torch.argmax(logits, dim=-1).cpu().numpy()
+        
+        # Map indices back to the alphabet characters
+        predicted_seq = "".join([AA_ALPHABET[idx] for idx in predicted_indices])
+        
+        return f"Predicted Sequence ({num_residues} residues):\n\n{predicted_seq}"
+        
+    except Exception as e:
+        return f"Error processing PDB: {str(e)}"
+
 
-# 3. Build the Gradio UI
+# ---------------------------------------------------------
+# 3. GRADIO UI
+# ---------------------------------------------------------
 demo = gr.Interface(
     fn=predict_sequence,
-    inputs=gr.File(label="Upload Protein Backbone (.pdb)", file_types=[".pdb"]),
-    outputs=gr.Textbox(label="Predicted Amino Acid Sequence", show_copy_button=True),
+    inputs=gr.File(label="Upload Target Protein Backbone (.pdb)", file_types=[".pdb"]),
+    outputs=gr.Textbox(label="Designed Amino Acid Sequence", show_copy_button=True, lines=5),
     title="Struct2Seq-GNN: Inverse Protein Folding",
-    description="Upload a 3D target backbone to generate a sequence optimized by a custom Graph Neural Network.",
-    allow_flagging="never"
+    description=(
+        "Upload a 3D target backbone to generate a sequence optimized by a custom Graph Neural Network.\n\n"
+        "**Model Performance:** Achieves ~30.3% global sequence recovery and **35.1% binding-pocket recovery** "
+        "on noisy coordinates, confirming strong generalization to underlying biophysical folding constraints."
+    ),
+    allow_flagging="never",
+    theme=gr.themes.Soft() # Adds a cleaner, more professional look
 )
 
 if __name__ == "__main__":