Update app.py

app.py

@@ -16,45 +16,39 @@ 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"
-    seed = 42
-
-
-# Assuming 'predict_stability' is your function that predicts protein stability
+    def __init__(self):
+        self.batch_size = 2
+        self.use_amp = False
+        self.num_workers = 1
+        self.max_length = 512
+        self.used_sequence = "left"
+        self.padding_side = "right"
+        self.task = "classification"
+        self.sequence_col = "sequence"
+        self.seed = 42
+
+
 def predict_stability(model_choice, organism_choice, pdb_file=None, sequence=None, cfg=Config()):
-    # 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")
         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:
-        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 output
-    else:
+        sequence = sequences[2] if model_choice == "SaProt" else sequences[0]
+
+    if not sequence:
         return "No valid input provided."
 
+    cell_line = "HeLa" if organism_choice == "Human" else "NIH3T3"
+    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 output
+
+
 
 def get_foldseek_seq(pdb_path):
     parsed_seqs = get_struc_seq(
@@ -93,7 +87,6 @@ def predict(cfg, sequence):
     model = PLTNUM_PreTrainedModel.from_pretrained(cfg.model_path, cfg=cfg)
     model.to(cfg.device)
 
-    # predictions = predict_fn(loader, model, cfg)
     model.eval()
     predictions = []
 
@@ -107,24 +100,23 @@ def predict(cfg, sequence):
                     else model(inputs)
                 )
         predictions += preds.cpu().tolist()
-    outputs = {}
+
     predictions = list(itertools.chain.from_iterable(predictions))
-    outputs["raw prediction values"] = predictions
-    outputs["binary prediction values"] = [1 if x > 0.5 else 0 for x in predictions]
+    outputs = {
+        "raw prediction values": predictions,
+        "binary prediction values": [1 if x > 0.5 else 0 for x in predictions]
+    }
 
     html_output = f"""
     <div style='border: 2px solid #4CAF50; padding: 10px; border-radius: 10px;'>
-        <h3 style='color: #4CAF50;'>Protein: {sequence}</h3>
-        <p><strong>Stability:</strong> {outputs['raw prediction values']}</p>
-        <p><strong>Organism:</strong> {outputs['binary prediction values']}</p>
+        <p><strong>Stability:</strong> {outputs['raw prediction values'][0]}</p>
+        <p><strong>Organism:</strong> {outputs['binary prediction values'][0]}</p>
     </div>
     """
     
     return html_output
 
 
-
-
 # Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -156,7 +148,6 @@ with gr.Blocks() as demo:
         with gr.TabItem("Upload PDB File"):
             gr.Markdown("### Upload your PDB file:")
             pdb_file = gr.File(label="Upload PDB File")
-
             predict_button = gr.Button("Predict Stability")
             prediction_output = gr.HTML(
                 label="Stability Prediction"