Update

.ipynb_checkpoints/app-checkpoint.py

@@ -54,6 +54,12 @@ def fetch_pdb(pdb_id):
         return pdb_path
     return None
 
+
+def normalize_scores(scores):
+    min_score = np.min(scores)
+    max_score = np.max(scores)
+    return (scores - min_score) / (max_score - min_score) if max_score > min_score else scores
+
 # Extract sequence and predict binding scores
 def process_pdb(pdb_id, segment):
     pdb_path = fetch_pdb(pdb_id)
@@ -70,9 +76,11 @@ def process_pdb(pdb_id, segment):
     with torch.no_grad():
         outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()
 
-    scores = outputs[:, 1] - outputs[:, 0]
+    scores = expit(outputs[:, 1] - outputs[:, 0])
+    normalized_scores = normalize_scores(scores)
+    
     result_str = "\n".join([
-        f"{res.get_resname()} {res.id[1]} {sequence[i]} {scores[i]:.2f}" 
+        f"{res.get_resname()} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}" 
         for i, res in enumerate(chain)
     ])
     
@@ -86,10 +94,14 @@ with gr.Blocks() as demo:
     gr.Markdown("# Protein Binding Site Prediction")
 
     with gr.Row():
-        pdb_input = gr.Textbox(label="PDB ID")
-        segment_input = gr.Textbox(label="Segment (Chain ID)")
-        visualize_btn = gr.Button("Visualize")
-        prediction_btn = gr.Button("Predict")
+        pdb_input = gr.Textbox(value="2IWI",
+                label="PDB ID",
+                placeholder="Enter PDB ID here...")
+        segment_input = gr.Textbox(value="A",
+                label="Chain ID (Segment)",
+                placeholder="Enter Chain ID here...")
+        visualize_btn = gr.Button("Visualize Sructure")
+        prediction_btn = gr.Button("Predict Ligand Binding Site")
 
     molecule_output = Molecule3D(label="Protein Structure", reps=reps)
     predictions_output = gr.Textbox(label="Binding Site Predictions")
@@ -102,4 +114,15 @@ with gr.Blocks() as demo:
         outputs=[predictions_output, molecule_output, download_output]
     )
 
+    gr.Markdown("## Examples")
+    gr.Examples(
+        examples=[
+            ["2IWI"],
+            ["7RPZ"],
+            ["3TJN"]
+        ],
+        inputs=[pdb_input, segment_input], 
+        outputs=[predictions_output, molecule_output, download_output]
+    )
+
 demo.launch(share=True)

app.py

@@ -54,6 +54,12 @@ def fetch_pdb(pdb_id):
         return pdb_path
     return None
 
+
+def normalize_scores(scores):
+    min_score = np.min(scores)
+    max_score = np.max(scores)
+    return (scores - min_score) / (max_score - min_score) if max_score > min_score else scores
+
 # Extract sequence and predict binding scores
 def process_pdb(pdb_id, segment):
     pdb_path = fetch_pdb(pdb_id)
@@ -70,9 +76,11 @@ def process_pdb(pdb_id, segment):
     with torch.no_grad():
         outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()
 
-    scores = outputs[:, 1] - outputs[:, 0]
+    scores = expit(outputs[:, 1] - outputs[:, 0])
+    normalized_scores = normalize_scores(scores)
+    
     result_str = "\n".join([
-        f"{res.get_resname()} {res.id[1]} {sequence[i]} {scores[i]:.2f}" 
+        f"{res.get_resname()} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}" 
         for i, res in enumerate(chain)
     ])
     
@@ -86,10 +94,14 @@ with gr.Blocks() as demo:
     gr.Markdown("# Protein Binding Site Prediction")
 
     with gr.Row():
-        pdb_input = gr.Textbox(label="PDB ID")
-        segment_input = gr.Textbox(label="Segment (Chain ID)")
-        visualize_btn = gr.Button("Visualize")
-        prediction_btn = gr.Button("Predict")
+        pdb_input = gr.Textbox(value="2IWI",
+                label="PDB ID",
+                placeholder="Enter PDB ID here...")
+        segment_input = gr.Textbox(value="A",
+                label="Chain ID (Segment)",
+                placeholder="Enter Chain ID here...")
+        visualize_btn = gr.Button("Visualize Sructure")
+        prediction_btn = gr.Button("Predict Ligand Binding Site")
 
     molecule_output = Molecule3D(label="Protein Structure", reps=reps)
     predictions_output = gr.Textbox(label="Binding Site Predictions")
@@ -102,4 +114,15 @@ with gr.Blocks() as demo:
         outputs=[predictions_output, molecule_output, download_output]
     )
 
+    gr.Markdown("## Examples")
+    gr.Examples(
+        examples=[
+            ["2IWI"],
+            ["7RPZ"],
+            ["3TJN"]
+        ],
+        inputs=[pdb_input, segment_input], 
+        outputs=[predictions_output, molecule_output, download_output]
+    )
+
 demo.launch(share=True)