Update app.py

app.py

@@ -2,23 +2,51 @@ import gradio as gr
 import sys
 import random
 import os
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
 sys.path.append("scripts/")
 from foldseek_util import get_struc_seq
+from utils import seed_everything
+from models import PLTNUM_PreTrainedModel
+from datasets import PLTNUMDataset
+
+class Config:
+    batch_size = 2
+    use_amp = False
+    num_workers = 1
+    max_length = 512
+    used_sequence = "left"
+    padding_side = "right"
+    task = "classification"
+    sequence_col = "sequence"
 
 # Assuming 'predict_stability' is your function that predicts protein stability
-def predict_stability(model_choice, organism_choice, pdb_file=None, sequence=None):
+def predict_stability(cfg, model_choice, organism_choice, pdb_file=None, sequence=None):
     # Check if pdb_file is provided
     if pdb_file:
         pdb_path = pdb_file.name  # Get the path of the uploaded PDB file
         os.system("chmod 777 bin/foldseek")
-        sequence = get_foldseek_seq(pdb_path)
-        if not sequence:
+        sequences = get_foldseek_seq(pdb_path)
+        if not sequences:
             return "Failed to extract sequence from the PDB file."
-    
+        if model_choice == "SaProt":
+            sequence = sequences[2]
+        else:
+            sequence = sequences[0]
+
+    if organism_choice == "Human":
+        cell_line = "HeLa"
+    else:
+        cell_line = "NIH3T3"
     # If sequence is provided directly
     if sequence:
-        # Add logic to predict stability using the sequence
-        return f"Predicted Stability using {model_choice} for {organism_choice}: Example Output with sequence {sequence[:10]}..."
+        cfg.model = f"sagawa/PLTNUM-{model_choice}-{cell_line}"
+        cfg.architecture = model_choice
+        cfg.model_path = f"sagawa/PLTNUM-{model_choice}-{cell_line}"
+        output = predict(cfg, sequence)
+        return f"Predicted Stability using {model_choice} for {organism_choice}: Example Output with sequence {sequence}..."
     else:
         return "No valid input provided."
     
@@ -33,6 +61,56 @@ def get_foldseek_seq(pdb_path):
     return parsed_seqs
 
 
+def predict(cfg, sequence):
+    cfg.token_length = 2 if cfg.architecture == "SaProt" else 1
+    cfg.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if cfg.used_sequence == "both":
+        cfg.max_length += 1
+
+    seed_everything(cfg.seed)
+
+    df = pd.DataFrame({cfg.sequence_col: [sequence]})
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        cfg.model_path, padding_side=cfg.padding_side
+    )
+    cfg.tokenizer = tokenizer
+
+    dataset = PLTNUMDataset(cfg, df, train=False)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=cfg.batch_size,
+        shuffle=False,
+        num_workers=cfg.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    model = PLTNUM_PreTrainedModel.from_pretrained(cfg.model_path, cfg=cfg)
+    model.to(cfg.device)
+
+    # predictions = predict_fn(loader, model, cfg)
+    model.eval()
+    predictions = []
+
+    for inputs, _ in dataloader:
+        inputs = inputs.to(cfg.device)
+        with torch.no_grad():
+            with torch.amp.autocast(enabled=cfg.use_amp):
+                preds = (
+                    torch.sigmoid(model(inputs))
+                    if cfg.task == "classification"
+                    else model(inputs)
+                )
+        predictions += preds.cpu().tolist()
+    outputs = {}
+    outputs["raw prediction values"] = predictions
+    outputs["binary prediction values"] = [1 if x > 0.5 else 0 for x in predictions]
+    return outputs   
+
+
+
 # Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -47,7 +125,7 @@ with gr.Blocks() as demo:
     # Model and Organism selection in the same row to avoid layout issues
     with gr.Row():
         model_choice = gr.Radio(
-            choices=["SaProt", "ESM-2"], 
+            choices=["SaProt", "ESM2"], 
             label="Select PLTNUM's base model.", 
             value="SaProt"
         )
@@ -82,7 +160,7 @@ with gr.Blocks() as demo:
     gr.Markdown(
         """
         ### How to Use:
-        - **Select Model**: Choose between 'SaProt' or 'ESM-2' for your prediction.
+        - **Select Model**: Choose between 'SaProt' or 'ESM2' for your prediction.
         - **Select Organism**: Choose between 'Mouse' or 'Human'.
         - **Upload PDB File**: Choose the 'Upload PDB File' tab and upload your file.
         - **Enter Sequence**: Alternatively, switch to the 'Enter Protein Sequence' tab and input your sequence.