Update

.ipynb_checkpoints/app-checkpoint.py

@@ -29,6 +29,22 @@ from datasets import Dataset
 
 from scipy.special import expit
 
+from datetime import datetime
+import gradio as gr
+import requests
+from Bio.PDB import PDBParser, MMCIFParser, PDBIO
+from Bio.PDB.Polypeptide import is_aa
+from Bio.SeqUtils import seq1
+from typing import Optional, Tuple
+import numpy as np
+import os
+from gradio_molecule3d import Molecule3D
+
+import re
+import pandas as pd
+import copy
+
+from scipy.special import expit
 
 
 # Load model and move to device
@@ -39,6 +55,24 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)
 model.eval()
 
+from datetime import datetime
+import gradio as gr
+import requests
+from Bio.PDB import PDBParser, MMCIFParser, PDBIO
+from Bio.PDB.Polypeptide import is_aa
+from Bio.SeqUtils import seq1
+from Bio.PDB import Select
+from typing import Optional, Tuple
+import numpy as np
+import os
+from gradio_molecule3d import Molecule3D
+
+import re
+import pandas as pd
+import copy
+
+from scipy.special import expit
+
 def normalize_scores(scores):
     min_score = np.min(scores)
     max_score = np.max(scores)
@@ -101,36 +135,44 @@ def fetch_pdb(pdb_id):
         pdb_path = convert_cif_to_pdb(pdb_path)
     return pdb_path
 
-def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list) -> str:
+def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list, protein_residues: list) -> str:
     """
-    Create a PDB file with only the specified chain and replace B-factor with prediction scores
+    Create a PDB file with only the selected chain and residues, replacing B-factor with prediction scores
     """
     # Read the original PDB file
     parser = PDBParser(QUIET=True)
     structure = parser.get_structure('protein', input_pdb)
     
-    # Prepare a new structure with only the specified chain
-    new_structure = structure.copy()
-    for model in new_structure:
-        # Remove all chains except the specified one
-        chains_to_remove = [chain for chain in model if chain.id != chain_id]
-        for chain in chains_to_remove:
-            model.detach_child(chain.id)
+    # Prepare a new structure with only the specified chain and selected residues
+    output_pdb = f"{os.path.splitext(input_pdb)[0]}_{chain_id}_predictions_scores.pdb"
     
-    # Create a modified PDB with scores in B-factor
+    # Create scores dictionary for easy lookup
     scores_dict = {resi: score for resi, score in residue_scores}
-    for model in new_structure:
-        for chain in model:
-            for residue in chain:
-                if residue.id[1] in scores_dict:
-                    for atom in residue:
-                        atom.bfactor = scores_dict[residue.id[1]] #* 100  # Scale score to B-factor range
-    
-    # Save the modified structure
-    output_pdb = f"{os.path.splitext(input_pdb)[0]}_{chain_id}_scored.pdb"
+
+    # Create a custom Select class
+    class ResidueSelector(Select):
+        def __init__(self, chain_id, selected_residues, scores_dict):
+            self.chain_id = chain_id
+            self.selected_residues = selected_residues
+            self.scores_dict = scores_dict
+        
+        def accept_chain(self, chain):
+            return chain.id == self.chain_id
+        
+        def accept_residue(self, residue):
+            return residue.id[1] in self.selected_residues
+
+        def accept_atom(self, atom):
+            if atom.parent.id[1] in self.scores_dict:
+                atom.bfactor = self.scores_dict[atom.parent.id[1]] * 100
+            return True
+
+    # Prepare output PDB with selected chain and residues, modified B-factors
     io = PDBIO()
-    io.set_structure(new_structure)
-    io.save(output_pdb)
+    selector = ResidueSelector(chain_id, [res.id[1] for res in protein_residues], scores_dict)
+    
+    io.set_structure(structure[0])
+    io.save(output_pdb, selector)
     
     return output_pdb
 
@@ -158,8 +200,6 @@ def calculate_geometric_center(pdb_path: str, high_score_residues: list, chain_i
         return center
     return None
 
-
-
 def process_pdb(pdb_id_or_file, segment):
     # Determine if input is a PDB ID or file path
     if pdb_id_or_file.endswith('.pdb'):
@@ -192,67 +232,75 @@ def process_pdb(pdb_id_or_file, segment):
     sequence = "".join(seq1(res.resname) for res in protein_residues)
     sequence_id = [res.id[1] for res in protein_residues]
     
-    # Prepare input for model prediction
-    input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
-    with torch.no_grad():
-        outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()
-    
-    # Calculate scores and normalize them
-    scores = expit(outputs[:, 1] - outputs[:, 0])
+    scores = np.random.rand(len(sequence))
     normalized_scores = normalize_scores(scores)
     
     # Zip residues with scores to track the residue ID and score
     residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]
 
-    # Identify high and mid scoring residues
-    high_score_residues = [resi for resi, score in residue_scores if score > 0.75]
-    mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75]
-
-    # Calculate geometric center of high-scoring residues
-    geo_center = calculate_geometric_center(pdb_path, high_score_residues, segment)
-    pymol_selection = f"select high_score_residues, resi {'+'.join(map(str, high_score_residues))} and chain {segment}"
-    pymol_center_cmd = f"show spheres, resi {'+'.join(map(str, high_score_residues))} and chain {segment}" if geo_center is not None else ""
-
-    # Generate the result string
+    
+    # Identify high scoring residues (> 0.5)
+    high_score_residues = [resi for resi, score in residue_scores if score > 0.5]
+    
+    # Preparing the result: only print high scoring residues
     current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     result_str = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\n\n"
+    result_str += "High-scoring Residues (Score > 0.5):\n"
     result_str += "Columns: Residue Name, Residue Number, One-letter Code, Normalized Score\n\n"
     result_str += "\n".join([
         f"{res.resname} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}" 
-        for i, res in enumerate(protein_residues)])
+        for i, res in enumerate(protein_residues) if res.id[1] in high_score_residues
+    ])
+
+    # Create chain-specific PDB with scores in B-factor
+    scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores, protein_residues)
+
+    # Molecule visualization with updated script with color mapping
+    mol_vis = molecule(pdb_path, residue_scores, segment)#, color_map)
+
+    # Improved PyMOL command suggestions
+    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    pymol_commands = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\n\n"
+    
+    pymol_commands += f"""
+# PyMOL Visualization Commands
+load {os.path.abspath(pdb_path)}, protein
+hide everything, all
+show cartoon, chain {segment}
+color white, chain {segment}
+"""
     
+    # Color specific residues
+    for score_range, color in [
+        (high_score_residues, "red")
+    ]:
+        if score_range:
+            resi_list = '+'.join(map(str, score_range))
+            pymol_commands += f"""
+select high_score_residues, resi {resi_list} and chain {segment}
+show sticks, high_score_residues
+color {color}, high_score_residues
+"""
     # Create prediction and scored PDB files
-    prediction_file = f"{pdb_id}_predictions.txt"
+    prediction_file = f"{pdb_id}_binding_site_residues.txt"
     with open(prediction_file, "w") as f:
         f.write(result_str)
-
-    # Create chain-specific PDB with scores in B-factor
-    scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores)
-
-    # Molecule visualization with updated script
-    mol_vis = molecule(pdb_path, residue_scores, segment)
-
-    # Construct PyMOL command suggestions
-    pymol_commands = f"""
-PyMOL Visualization Commands:
-1. Load PDB: load {os.path.abspath(pdb_path)}
-2. Select high-scoring residues: {pymol_selection}
-3. Highlight high-scoring residues: show sticks, high_score_residues
-{pymol_center_cmd}
-"""
     
-    return result_str + "\n\n" + pymol_commands, mol_vis, [prediction_file, scored_pdb]
-
+    return pymol_commands, mol_vis, [prediction_file,scored_pdb]
 
 def molecule(input_pdb, residue_scores=None, segment='A'):
+    # More granular scoring for visualization
     mol = read_mol(input_pdb)  # Read PDB file content
 
     # Prepare high-scoring residues script if scores are provided
     high_score_script = ""
     if residue_scores is not None:
         # Filter residues based on their scores
-        high_score_residues = [resi for resi, score in residue_scores if score > 0.75]
-        mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75]
+        class1_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.6]
+        class2_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.7]
+        class3_score_residues = [resi for resi, score in residue_scores if 0.7 < score <= 0.8]
+        class4_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 0.9]
+        class5_score_residues = [resi for resi, score in residue_scores if 0.9 < score <= 1.0]
         
         high_score_script = """
         // Load the original model and apply white cartoon style
@@ -264,26 +312,57 @@ def molecule(input_pdb, residue_scores=None, segment='A'):
         );
 
         // Create a new model for high-scoring residues and apply red sticks style
-        let highScoreModel = viewer.addModel(pdb, "pdb");
-        highScoreModel.setStyle({}, {});
-        highScoreModel.setStyle(
+        let class1Model = viewer.addModel(pdb, "pdb");
+        class1Model.setStyle({}, {});
+        class1Model.setStyle(
             {"chain": "%s", "resi": [%s]}, 
-            {"stick": {"color": "red"}}
+            {"stick": {"color": "blue"}}
         );
 
-        // Create a new model for medium-scoring residues and apply orange sticks style
-        let midScoreModel = viewer.addModel(pdb, "pdb");
-        midScoreModel.setStyle({}, {});
-        midScoreModel.setStyle(
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class2Model = viewer.addModel(pdb, "pdb");
+        class2Model.setStyle({}, {});
+        class2Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "lightblue"}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class3Model = viewer.addModel(pdb, "pdb");
+        class3Model.setStyle({}, {});
+        class3Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "white"}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class4Model = viewer.addModel(pdb, "pdb");
+        class4Model.setStyle({}, {});
+        class4Model.setStyle(
             {"chain": "%s", "resi": [%s]}, 
             {"stick": {"color": "orange"}}
         );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class5Model = viewer.addModel(pdb, "pdb");
+        class5Model.setStyle({}, {});
+        class5Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "red"}}
+        );
+
         """ % (
             segment,
             segment,
-            ", ".join(str(resi) for resi in high_score_residues),
+            ", ".join(str(resi) for resi in class1_score_residues),
             segment,
-            ", ".join(str(resi) for resi in mid_score_residues)
+            ", ".join(str(resi) for resi in class2_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class3_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class4_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class5_score_residues)
         )
     
     # Generate the full HTML content
@@ -351,14 +430,22 @@ def molecule(input_pdb, residue_scores=None, segment='A'):
     # Return the HTML content within an iframe safely encoded for special characters
     return f'<iframe width="100%" height="700" srcdoc="{html_content.replace(chr(34), "&quot;").replace(chr(39), "&#39;")}"></iframe>'
 
-
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# Protein Binding Site Prediction")
     
-    with gr.Row():
-        pdb_input = gr.Textbox(value="4BDU", label="PDB ID", placeholder="Enter PDB ID here...")
-        visualize_btn = gr.Button("Visualize Structure")
+    # Mode selection
+    mode = gr.Radio(
+        choices=["PDB ID", "Upload File"],
+        value="PDB ID",
+        label="Input Mode",
+        info="Choose whether to input a PDB ID or upload a PDB/CIF file."
+    )
+
+    # Input components based on mode
+    pdb_input = gr.Textbox(value="4BDU", label="PDB ID", placeholder="Enter PDB ID here...")
+    pdb_file = gr.File(label="Upload PDB/CIF File", visible=False)
+    visualize_btn = gr.Button("Visualize Structure")
 
     molecule_output2 = Molecule3D(label="Protein Structure", reps=[
         {
@@ -375,23 +462,70 @@ with gr.Blocks() as demo:
         segment_input = gr.Textbox(value="A", label="Chain ID", placeholder="Enter Chain ID here...")
         prediction_btn = gr.Button("Predict Binding Site")
 
-
     molecule_output = gr.HTML(label="Protein Structure")
-    predictions_output = gr.Textbox(label="Binding Site Predictions")
+    explanation_vis = gr.Markdown("""
+    Residues with a score > 0.5 are considered binding sites and represented as sticks with the score dependent colorcoding:
+    - 0.5-0.6: blue  
+    - 0.6–0.7: light blue  
+    - 0.7–0.8: white
+    - 0.8–0.9: orange
+    - 0.9–1.0: red
+    """)
+    predictions_output = gr.Textbox(label="Visualize Prediction with PyMol")
+    gr.Markdown("### Download:\n- List of predicted binding site residues\n- PDB with score in beta factor column")
     download_output = gr.File(label="Download Files", file_count="multiple")
     
+    def process_interface(mode, pdb_id, pdb_file, chain_id):
+        if mode == "PDB ID":
+            return process_pdb(pdb_id, chain_id)
+        elif mode == "Upload File":
+            _, ext = os.path.splitext(pdb_file.name)
+            file_path = os.path.join('./', f"{_}{ext}")
+            if ext == '.cif':
+                pdb_path = convert_cif_to_pdb(file_path)
+            else:
+                pdb_path= file_path
+            return process_pdb(pdb_path, chain_id)
+        else:
+            return "Error: Invalid mode selected", None, None
+
+    def fetch_interface(mode, pdb_id, pdb_file):
+        if mode == "PDB ID":
+            return fetch_pdb(pdb_id)
+        elif mode == "Upload File":
+            _, ext = os.path.splitext(pdb_file.name)
+            file_path = os.path.join('./', f"{_}{ext}")
+            #print(ext)
+            if ext == '.cif':
+                pdb_path = convert_cif_to_pdb(file_path)
+            else:
+                pdb_path= file_path
+            #print(pdb_path)
+            return pdb_path
+        else:
+            return "Error: Invalid mode selected"
+
+    def toggle_mode(selected_mode):
+        if selected_mode == "PDB ID":
+            return gr.update(visible=True), gr.update(visible=False)
+        else:
+            return gr.update(visible=False), gr.update(visible=True)
+
+    mode.change(
+        toggle_mode,
+        inputs=[mode],
+        outputs=[pdb_input, pdb_file]
+    )
+
     prediction_btn.click(
-        process_pdb, 
-        inputs=[
-            pdb_input, 
-            segment_input
-        ], 
+        process_interface, 
+        inputs=[mode, pdb_input, pdb_file, segment_input], 
         outputs=[predictions_output, molecule_output, download_output]
     )
 
     visualize_btn.click(
-        fetch_pdb, 
-        inputs=[pdb_input], 
+        fetch_interface, 
+        inputs=[mode, pdb_input, pdb_file], 
         outputs=molecule_output2
     )
 

.ipynb_checkpoints/test3-checkpoint.ipynb

@@ -1149,7 +1149,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 38,
+   "execution_count": 39,
    "id": "514fad12-a31a-495f-af9e-04a18e11175e",
    "metadata": {},
    "outputs": [
@@ -1157,8 +1157,8 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "* Running on local URL:  http://127.0.0.1:7896\n",
-      "* Running on public URL: https://387fb4706015321f92.gradio.live\n",
+      "* Running on local URL:  http://127.0.0.1:7897\n",
+      "* Running on public URL: https://0d9b5d36fa5302e0df.gradio.live\n",
       "\n",
       "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
      ]
@@ -1166,7 +1166,7 @@
     {
      "data": {
       "text/html": [
-       "<div><iframe src=\"https://387fb4706015321f92.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+       "<div><iframe src=\"https://0d9b5d36fa5302e0df.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"
@@ -1179,7 +1179,7 @@
      "data": {
       "text/plain": []
      },
-     "execution_count": 38,
+     "execution_count": 39,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -1422,7 +1422,7 @@
     "\n",
     "        // Create a new model for medium-scoring residues and apply orange sticks style\n",
     "        let midScoreModel = viewer.addModel(pdb, \"pdb\");\n",
-    "        midScoreModel.setStyle({}, {});\n",
+    "        highScormidScoreModeleModel.setStyle({}, {});\n",
     "        midScoreModel.setStyle(\n",
     "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
     "            {\"stick\": {\"color\": \"orange\"}}\n",

2IWI_predictions.txt

@@ -1,249 +0,0 @@
-GLY 22 G 0.18
-LYS 23 K 0.51
-ASP 24 D 0.12
-ARG 25 R 0.25
-GLU 26 E 0.08
-ALA 27 A 0.82
-PHE 28 F 0.65
-GLU 29 E 0.65
-ALA 30 A 0.22
-GLU 31 E 0.49
-TYR 32 Y 0.57
-ARG 33 R 0.56
-LEU 34 L 0.83
-GLY 35 G 0.42
-PRO 36 P 0.97
-LEU 37 L 0.65
-LEU 38 L 0.08
-GLY 39 G 0.05
-LYS 40 K 0.55
-GLY 41 G 0.38
-GLY 42 G 0.45
-PHE 43 F 0.92
-GLY 44 G 0.00
-THR 45 T 0.76
-VAL 46 V 0.63
-PHE 47 F 0.97
-ALA 48 A 0.57
-GLY 49 G 0.94
-HIS 50 H 0.40
-ARG 51 R 0.27
-LEU 52 L 0.65
-THR 53 T 0.84
-ASP 54 D 0.85
-ARG 55 R 0.46
-LEU 56 L 0.87
-GLN 57 Q 0.76
-VAL 58 V 0.22
-ALA 59 A 0.65
-ILE 60 I 0.87
-LYS 61 K 0.69
-VAL 62 V 0.76
-ILE 63 I 0.70
-PRO 64 P 0.04
-ARG 65 R 0.20
-THR 79 T 0.80
-CYS 80 C 0.82
-PRO 81 P 0.72
-LEU 82 L 0.17
-GLU 83 E 0.70
-VAL 84 V 0.21
-ALA 85 A 0.15
-LEU 86 L 0.28
-LEU 87 L 0.03
-TRP 88 W 0.18
-LYS 89 K 0.01
-VAL 90 V 0.43
-GLY 91 G 0.25
-ALA 92 A 0.65
-GLY 93 G 0.00
-GLY 94 G 0.52
-GLY 95 G 0.22
-HIS 96 H 0.03
-PRO 97 P 0.57
-GLY 98 G 0.32
-VAL 99 V 0.89
-ILE 100 I 0.14
-ARG 101 R 0.66
-LEU 102 L 0.18
-LEU 103 L 0.30
-ASP 104 D 0.36
-TRP 105 W 0.83
-PHE 106 F 0.77
-GLU 107 E 0.95
-PHE 112 F 0.04
-MET 113 M 0.05
-LEU 114 L 0.32
-VAL 115 V 1.00
-LEU 116 L 0.43
-GLU 117 E 0.76
-ARG 118 R 0.65
-PRO 119 P 0.28
-LEU 120 L 0.74
-PRO 121 P 0.69
-ALA 122 A 0.89
-GLN 123 Q 0.68
-ASP 124 D 0.67
-LEU 125 L 0.89
-PHE 126 F 0.33
-ASP 127 D 0.05
-TYR 128 Y 0.59
-ILE 129 I 0.19
-THR 130 T 0.88
-GLU 131 E 0.24
-LYS 132 K 0.04
-GLY 133 G 0.99
-PRO 134 P 0.43
-LEU 135 L 0.31
-GLY 136 G 0.83
-GLU 137 E 0.12
-GLY 138 G 0.02
-PRO 139 P 0.71
-SER 140 S 0.70
-ARG 141 R 0.63
-CYS 142 C 0.70
-PHE 143 F 0.92
-PHE 144 F 0.02
-GLY 145 G 0.72
-GLN 146 Q 0.03
-VAL 147 V 0.70
-VAL 148 V 0.34
-ALA 149 A 0.95
-ALA 150 A 0.39
-ILE 151 I 0.21
-GLN 152 Q 0.86
-HIS 153 H 0.11
-CYS 154 C 0.30
-HIS 155 H 0.12
-SER 156 S 0.55
-ARG 157 R 0.20
-GLY 158 G 0.32
-VAL 159 V 0.80
-VAL 160 V 0.43
-HIS 161 H 0.99
-ARG 162 R 0.13
-ASP 163 D 0.73
-ILE 164 I 0.70
-LYS 165 K 0.88
-ASP 166 D 0.56
-GLU 167 E 0.61
-ASN 168 N 0.01
-ILE 169 I 0.48
-LEU 170 L 0.18
-ILE 171 I 0.28
-ASP 172 D 0.79
-LEU 173 L 0.33
-ARG 174 R 0.31
-ARG 175 R 0.39
-GLY 176 G 0.19
-CYS 177 C 0.57
-ALA 178 A 0.99
-LYS 179 K 0.47
-LEU 180 L 0.02
-ILE 181 I 0.81
-ASP 182 D 0.59
-PHE 183 F 0.74
-GLY 184 G 0.43
-SER 185 S 0.90
-GLY 186 G 0.87
-ALA 187 A 0.39
-LEU 188 L 0.43
-LEU 189 L 0.84
-HIS 190 H 0.91
-ASP 191 D 0.45
-GLU 192 E 0.00
-PRO 193 P 0.86
-TYR 194 Y 0.11
-THR 195 T 0.54
-ASP 196 D 0.70
-PHE 197 F 0.62
-ASP 198 D 0.31
-GLY 199 G 0.41
-THR 200 T 0.85
-ARG 201 R 0.18
-VAL 202 V 0.10
-TYR 203 Y 0.22
-SER 204 S 0.31
-PRO 205 P 0.41
-PRO 206 P 0.87
-GLU 207 E 0.77
-TRP 208 W 0.51
-ILE 209 I 0.18
-SER 210 S 0.03
-ARG 211 R 0.41
-HIS 212 H 0.83
-GLN 213 Q 0.30
-TYR 214 Y 0.38
-HIS 215 H 0.28
-ALA 216 A 0.51
-LEU 217 L 0.61
-PRO 218 P 0.77
-ALA 219 A 0.79
-THR 220 T 0.32
-VAL 221 V 0.35
-TRP 222 W 0.44
-SER 223 S 0.35
-LEU 224 L 0.67
-GLY 225 G 0.21
-ILE 226 I 0.88
-LEU 227 L 0.38
-LEU 228 L 0.27
-TYR 229 Y 0.53
-ASP 230 D 0.36
-MET 231 M 0.76
-VAL 232 V 0.59
-CYS 233 C 0.44
-GLY 234 G 0.88
-ASP 235 D 0.54
-ILE 236 I 0.63
-PRO 237 P 0.41
-PHE 238 F 0.84
-GLU 239 E 0.66
-ARG 240 R 0.20
-ASP 241 D 0.08
-GLN 242 Q 0.23
-GLU 243 E 0.31
-ILE 244 I 0.17
-LEU 245 L 0.58
-GLU 246 E 0.76
-ALA 247 A 0.82
-GLU 248 E 0.39
-LEU 249 L 0.53
-HIS 250 H 0.67
-PHE 251 F 0.36
-PRO 252 P 0.16
-ALA 253 A 0.08
-HIS 254 H 0.53
-VAL 255 V 0.39
-SER 256 S 0.24
-PRO 257 P 0.06
-ASP 258 D 0.79
-CYS 259 C 0.54
-CYS 260 C 0.46
-ALA 261 A 0.29
-LEU 262 L 0.60
-ILE 263 I 0.33
-ARG 264 R 0.56
-ARG 265 R 0.95
-CYS 266 C 0.63
-LEU 267 L 0.83
-ALA 268 A 0.22
-PRO 269 P 0.18
-LYS 270 K 0.71
-PRO 271 P 0.91
-SER 272 S 0.84
-SER 273 S 0.62
-ARG 274 R 0.22
-PRO 275 P 0.34
-SER 276 S 0.74
-LEU 277 L 0.41
-GLU 278 E 0.78
-GLU 279 E 0.76
-ILE 280 I 0.40
-LEU 281 L 0.27
-LEU 282 L 0.23
-ASP 283 D 0.65
-PRO 284 P 0.45
-TRP 285 W 0.72
-MET 286 M 0.57
-GLN 287 Q 0.29

4BDU_predictions.txt

@@ -1,300 +0,0 @@
-Prediction for PDB: 4BDU, Chain: A
-Date: 2024-12-11 16:57:50
-
-Columns: Residue Name, Residue Number, One-letter Code, Normalized Score
-
-SER 2 S 0.05
-LYS 3 K 0.39
-GLY 4 G 0.24
-GLU 5 E 0.26
-GLU 6 E 0.35
-LEU 7 L 0.45
-PHE 8 F 0.82
-THR 9 T 0.32
-GLY 10 G 0.73
-VAL 11 V 0.42
-VAL 12 V 0.33
-PRO 13 P 0.96
-ILE 14 I 0.68
-LEU 15 L 0.71
-VAL 16 V 0.84
-GLU 17 E 0.26
-LEU 18 L 0.54
-ASP 19 D 0.46
-GLY 20 G 0.12
-ASP 21 D 0.57
-VAL 22 V 0.32
-ASN 23 N 0.18
-GLY 24 G 0.48
-HIS 25 H 0.95
-LYS 26 K 0.88
-PHE 27 F 0.13
-SER 28 S 0.12
-VAL 29 V 0.58
-SER 30 S 0.19
-GLY 31 G 0.09
-GLU 32 E 0.17
-GLY 33 G 0.60
-GLU 34 E 0.92
-GLY 35 G 0.48
-ASP 36 D 0.35
-ALA 37 A 0.72
-THR 38 T 0.47
-TYR 39 Y 0.11
-GLY 40 G 0.57
-LYS 41 K 0.86
-LEU 42 L 0.42
-THR 43 T 0.98
-LEU 44 L 0.27
-LYS 45 K 0.05
-PHE 46 F 0.54
-ILE 47 I 0.25
-CYS 48 C 0.73
-THR 49 T 0.44
-THR 50 T 0.85
-GLY 51 G 0.17
-LYS 52 K 0.72
-LEU 53 L 0.03
-PRO 54 P 0.26
-VAL 55 V 0.64
-PRO 56 P 0.88
-TRP 57 W 0.84
-PRO 58 P 0.71
-THR 59 T 0.41
-LEU 60 L 0.18
-VAL 61 V 0.32
-THR 62 T 0.87
-THR 63 T 0.87
-PHE 64 F 1.00
-VAL 68 V 0.50
-GLN 69 Q 0.10
-CYS 70 C 0.71
-PHE 71 F 0.47
-SER 72 S 0.46
-ARG 73 R 0.99
-TYR 74 Y 0.40
-PRO 75 P 0.78
-ASP 76 D 0.42
-HIS 77 H 0.93
-MET 78 M 0.47
-LYS 79 K 0.51
-GLN 80 Q 0.85
-HIS 81 H 0.11
-ASP 82 D 0.87
-PHE 83 F 0.13
-PHE 84 F 0.56
-LYS 85 K 0.44
-SER 86 S 0.44
-ALA 87 A 0.20
-MET 88 M 0.33
-PRO 89 P 0.77
-GLU 90 E 0.32
-GLY 91 G 0.80
-TYR 92 Y 0.52
-VAL 93 V 0.46
-GLN 94 Q 0.26
-GLU 95 E 0.03
-ARG 96 R 0.99
-THR 97 T 0.72
-ILE 98 I 0.38
-PHE 99 F 0.63
-PHE 100 F 0.03
-LYS 101 K 0.10
-ASP 102 D 0.52
-ASP 103 D 0.41
-GLY 104 G 0.91
-ASN 105 N 0.17
-TYR 106 Y 0.75
-LYS 107 K 0.07
-THR 108 T 0.78
-ARG 109 R 0.21
-ALA 110 A 0.93
-GLU 111 E 0.34
-VAL 112 V 0.06
-LYS 113 K 0.92
-PHE 114 F 0.43
-GLU 115 E 0.22
-GLY 116 G 0.67
-ASP 117 D 0.54
-THR 118 T 0.18
-LEU 119 L 0.33
-VAL 120 V 0.52
-ASN 121 N 0.23
-ARG 122 R 0.18
-ILE 123 I 0.52
-GLU 124 E 0.85
-LEU 125 L 0.66
-LYS 126 K 0.69
-GLY 127 G 0.46
-ILE 128 I 0.48
-ASP 129 D 0.55
-PHE 130 F 0.90
-LYS 131 K 1.00
-GLU 132 E 0.98
-ASP 133 D 0.41
-GLY 134 G 0.78
-ASN 135 N 0.12
-ILE 136 I 0.06
-LEU 137 L 0.80
-GLY 138 G 0.70
-HIS 139 H 0.52
-LYS 140 K 0.40
-LEU 141 L 0.97
-GLU 142 E 0.25
-TYR 143 Y 0.53
-ASN 144 N 0.26
-TYR 145 Y 0.67
-ASN 146 N 0.65
-SER 147 S 0.91
-HIS 148 H 0.82
-ASN 149 N 0.93
-VAL 150 V 0.67
-TYR 151 Y 0.87
-ILE 152 I 0.02
-MET 153 M 0.37
-ALA 154 A 0.50
-ASP 155 D 0.89
-LYS 156 K 1.00
-GLN 157 Q 0.96
-LYS 158 K 0.83
-ASN 159 N 0.95
-GLY 160 G 0.02
-ILE 161 I 0.57
-LYS 162 K 0.82
-VAL 163 V 0.66
-ASN 164 N 0.32
-PHE 165 F 0.50
-LYS 166 K 0.11
-ILE 167 I 0.49
-ARG 168 R 0.20
-HIS 169 H 0.82
-ASN 170 N 0.34
-ILE 171 I 0.91
-GLU 172 E 0.28
-ASP 173 D 0.02
-GLY 174 G 0.09
-SER 175 S 0.44
-VAL 176 V 0.87
-GLN 177 Q 0.65
-LEU 178 L 0.88
-ALA 179 A 0.89
-ASP 180 D 0.53
-HIS 181 H 0.89
-TYR 182 Y 0.44
-GLN 183 Q 0.02
-GLN 184 Q 0.91
-ASN 185 N 0.57
-THR 186 T 0.00
-PRO 187 P 0.97
-ILE 188 I 0.17
-GLY 189 G 0.57
-ASP 190 D 0.46
-GLY 191 G 0.08
-PRO 192 P 0.85
-VAL 193 V 0.09
-LEU 194 L 0.79
-LEU 195 L 0.61
-PRO 196 P 0.72
-ASP 197 D 0.29
-ASN 198 N 0.95
-HIS 199 H 0.78
-TYR 200 Y 0.02
-LEU 201 L 0.55
-SER 202 S 0.63
-THR 203 T 0.38
-GLN 204 Q 0.18
-SER 205 S 0.48
-ASN 206 N 0.19
-LEU 207 L 0.71
-SER 208 S 0.56
-LYS 209 K 0.56
-ASP 210 D 0.98
-PRO 211 P 0.43
-ASN 212 N 0.91
-GLU 213 E 0.76
-LYS 214 K 0.58
-ARG 215 R 0.42
-ASP 216 D 0.81
-HIS 217 H 0.96
-MET 218 M 0.26
-VAL 219 V 0.01
-LEU 220 L 0.27
-LEU 221 L 0.26
-GLU 222 E 0.92
-PHE 223 F 0.84
-VAL 224 V 0.72
-THR 225 T 1.00
-ALA 226 A 0.55
-ALA 227 A 0.72
-GLY 228 G 0.44
-ILE 229 I 0.01
-THR 230 T 0.98
-ALA 1054 A 0.83
-SER 1055 S 0.78
-THR 1056 T 0.55
-LYS 1057 K 0.40
-LYS 1058 K 0.06
-LEU 1059 L 0.82
-SER 1060 S 0.59
-GLU 1061 E 0.68
-SER 1062 S 0.28
-LEU 1063 L 0.79
-LYS 1064 K 0.94
-ARG 1065 R 0.32
-ILE 1066 I 0.28
-GLY 1067 G 0.94
-ASP 1068 D 0.19
-GLU 1069 E 0.76
-LEU 1070 L 0.19
-ASP 1071 D 0.14
-SER 1072 S 0.04
-ASN 1073 N 0.39
-MET 1074 M 0.50
-GLU 1075 E 0.92
-LEU 1076 L 0.81
-GLN 1077 Q 0.04
-ARG 1078 R 0.97
-MET 1079 M 0.20
-ILE 1080 I 0.90
-ALA 1081 A 0.43
-ALA 1082 A 0.93
-VAL 1083 V 0.28
-ASP 1084 D 0.29
-THR 1085 T 0.83
-ASP 1086 D 0.79
-SER 1087 S 0.39
-PRO 1088 P 0.85
-ARG 1089 R 0.41
-GLU 1090 E 0.08
-VAL 1091 V 0.10
-PHE 1092 F 0.15
-PHE 1093 F 0.10
-ARG 1094 R 0.59
-VAL 1095 V 0.69
-ALA 1096 A 0.50
-ALA 1097 A 0.86
-ASP 1098 D 0.77
-MET 1099 M 0.60
-PHE 1100 F 0.13
-SER 1101 S 0.22
-ASP 1102 D 0.29
-GLY 1103 G 0.22
-ASN 1104 N 0.01
-PHE 1105 F 0.24
-ASN 1106 N 0.48
-TRP 1107 W 0.45
-GLY 1108 G 0.52
-ARG 1109 R 0.86
-VAL 1110 V 0.68
-VAL 1111 V 0.96
-ALA 1112 A 0.01
-LEU 1113 L 0.88
-PHE 1114 F 0.66
-TYR 1115 Y 0.11
-PHE 1116 F 0.62
-ALA 1117 A 0.62
-SER 1118 S 0.26
-LYS 1119 K 0.58
-LEU 1120 L 0.18
-VAL 1121 V 0.85
-LEU 1122 L 0.27

app.py

@@ -29,6 +29,22 @@ from datasets import Dataset
 
 from scipy.special import expit
 
+from datetime import datetime
+import gradio as gr
+import requests
+from Bio.PDB import PDBParser, MMCIFParser, PDBIO
+from Bio.PDB.Polypeptide import is_aa
+from Bio.SeqUtils import seq1
+from typing import Optional, Tuple
+import numpy as np
+import os
+from gradio_molecule3d import Molecule3D
+
+import re
+import pandas as pd
+import copy
+
+from scipy.special import expit
 
 
 # Load model and move to device
@@ -39,6 +55,24 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)
 model.eval()
 
+from datetime import datetime
+import gradio as gr
+import requests
+from Bio.PDB import PDBParser, MMCIFParser, PDBIO
+from Bio.PDB.Polypeptide import is_aa
+from Bio.SeqUtils import seq1
+from Bio.PDB import Select
+from typing import Optional, Tuple
+import numpy as np
+import os
+from gradio_molecule3d import Molecule3D
+
+import re
+import pandas as pd
+import copy
+
+from scipy.special import expit
+
 def normalize_scores(scores):
     min_score = np.min(scores)
     max_score = np.max(scores)
@@ -101,36 +135,44 @@ def fetch_pdb(pdb_id):
         pdb_path = convert_cif_to_pdb(pdb_path)
     return pdb_path
 
-def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list) -> str:
+def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list, protein_residues: list) -> str:
     """
-    Create a PDB file with only the specified chain and replace B-factor with prediction scores
+    Create a PDB file with only the selected chain and residues, replacing B-factor with prediction scores
     """
     # Read the original PDB file
     parser = PDBParser(QUIET=True)
     structure = parser.get_structure('protein', input_pdb)
     
-    # Prepare a new structure with only the specified chain
-    new_structure = structure.copy()
-    for model in new_structure:
-        # Remove all chains except the specified one
-        chains_to_remove = [chain for chain in model if chain.id != chain_id]
-        for chain in chains_to_remove:
-            model.detach_child(chain.id)
+    # Prepare a new structure with only the specified chain and selected residues
+    output_pdb = f"{os.path.splitext(input_pdb)[0]}_{chain_id}_predictions_scores.pdb"
     
-    # Create a modified PDB with scores in B-factor
+    # Create scores dictionary for easy lookup
     scores_dict = {resi: score for resi, score in residue_scores}
-    for model in new_structure:
-        for chain in model:
-            for residue in chain:
-                if residue.id[1] in scores_dict:
-                    for atom in residue:
-                        atom.bfactor = scores_dict[residue.id[1]] #* 100  # Scale score to B-factor range
-    
-    # Save the modified structure
-    output_pdb = f"{os.path.splitext(input_pdb)[0]}_{chain_id}_scored.pdb"
+
+    # Create a custom Select class
+    class ResidueSelector(Select):
+        def __init__(self, chain_id, selected_residues, scores_dict):
+            self.chain_id = chain_id
+            self.selected_residues = selected_residues
+            self.scores_dict = scores_dict
+        
+        def accept_chain(self, chain):
+            return chain.id == self.chain_id
+        
+        def accept_residue(self, residue):
+            return residue.id[1] in self.selected_residues
+
+        def accept_atom(self, atom):
+            if atom.parent.id[1] in self.scores_dict:
+                atom.bfactor = self.scores_dict[atom.parent.id[1]] * 100
+            return True
+
+    # Prepare output PDB with selected chain and residues, modified B-factors
     io = PDBIO()
-    io.set_structure(new_structure)
-    io.save(output_pdb)
+    selector = ResidueSelector(chain_id, [res.id[1] for res in protein_residues], scores_dict)
+    
+    io.set_structure(structure[0])
+    io.save(output_pdb, selector)
     
     return output_pdb
 
@@ -158,8 +200,6 @@ def calculate_geometric_center(pdb_path: str, high_score_residues: list, chain_i
         return center
     return None
 
-
-
 def process_pdb(pdb_id_or_file, segment):
     # Determine if input is a PDB ID or file path
     if pdb_id_or_file.endswith('.pdb'):
@@ -192,67 +232,75 @@ def process_pdb(pdb_id_or_file, segment):
     sequence = "".join(seq1(res.resname) for res in protein_residues)
     sequence_id = [res.id[1] for res in protein_residues]
     
-    # Prepare input for model prediction
-    input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
-    with torch.no_grad():
-        outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()
-    
-    # Calculate scores and normalize them
-    scores = expit(outputs[:, 1] - outputs[:, 0])
+    scores = np.random.rand(len(sequence))
     normalized_scores = normalize_scores(scores)
     
     # Zip residues with scores to track the residue ID and score
     residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]
 
-    # Identify high and mid scoring residues
-    high_score_residues = [resi for resi, score in residue_scores if score > 0.75]
-    mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75]
-
-    # Calculate geometric center of high-scoring residues
-    geo_center = calculate_geometric_center(pdb_path, high_score_residues, segment)
-    pymol_selection = f"select high_score_residues, resi {'+'.join(map(str, high_score_residues))} and chain {segment}"
-    pymol_center_cmd = f"show spheres, resi {'+'.join(map(str, high_score_residues))} and chain {segment}" if geo_center is not None else ""
-
-    # Generate the result string
+    
+    # Identify high scoring residues (> 0.5)
+    high_score_residues = [resi for resi, score in residue_scores if score > 0.5]
+    
+    # Preparing the result: only print high scoring residues
     current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     result_str = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\n\n"
+    result_str += "High-scoring Residues (Score > 0.5):\n"
     result_str += "Columns: Residue Name, Residue Number, One-letter Code, Normalized Score\n\n"
     result_str += "\n".join([
         f"{res.resname} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}" 
-        for i, res in enumerate(protein_residues)])
+        for i, res in enumerate(protein_residues) if res.id[1] in high_score_residues
+    ])
+
+    # Create chain-specific PDB with scores in B-factor
+    scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores, protein_residues)
+
+    # Molecule visualization with updated script with color mapping
+    mol_vis = molecule(pdb_path, residue_scores, segment)#, color_map)
+
+    # Improved PyMOL command suggestions
+    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    pymol_commands = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\n\n"
+    
+    pymol_commands += f"""
+# PyMOL Visualization Commands
+load {os.path.abspath(pdb_path)}, protein
+hide everything, all
+show cartoon, chain {segment}
+color white, chain {segment}
+"""
     
+    # Color specific residues
+    for score_range, color in [
+        (high_score_residues, "red")
+    ]:
+        if score_range:
+            resi_list = '+'.join(map(str, score_range))
+            pymol_commands += f"""
+select high_score_residues, resi {resi_list} and chain {segment}
+show sticks, high_score_residues
+color {color}, high_score_residues
+"""
     # Create prediction and scored PDB files
-    prediction_file = f"{pdb_id}_predictions.txt"
+    prediction_file = f"{pdb_id}_binding_site_residues.txt"
     with open(prediction_file, "w") as f:
         f.write(result_str)
-
-    # Create chain-specific PDB with scores in B-factor
-    scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores)
-
-    # Molecule visualization with updated script
-    mol_vis = molecule(pdb_path, residue_scores, segment)
-
-    # Construct PyMOL command suggestions
-    pymol_commands = f"""
-PyMOL Visualization Commands:
-1. Load PDB: load {os.path.abspath(pdb_path)}
-2. Select high-scoring residues: {pymol_selection}
-3. Highlight high-scoring residues: show sticks, high_score_residues
-{pymol_center_cmd}
-"""
     
-    return result_str + "\n\n" + pymol_commands, mol_vis, [prediction_file, scored_pdb]
-
+    return pymol_commands, mol_vis, [prediction_file,scored_pdb]
 
 def molecule(input_pdb, residue_scores=None, segment='A'):
+    # More granular scoring for visualization
     mol = read_mol(input_pdb)  # Read PDB file content
 
     # Prepare high-scoring residues script if scores are provided
     high_score_script = ""
     if residue_scores is not None:
         # Filter residues based on their scores
-        high_score_residues = [resi for resi, score in residue_scores if score > 0.75]
-        mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75]
+        class1_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.6]
+        class2_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.7]
+        class3_score_residues = [resi for resi, score in residue_scores if 0.7 < score <= 0.8]
+        class4_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 0.9]
+        class5_score_residues = [resi for resi, score in residue_scores if 0.9 < score <= 1.0]
         
         high_score_script = """
         // Load the original model and apply white cartoon style
@@ -264,26 +312,57 @@ def molecule(input_pdb, residue_scores=None, segment='A'):
         );
 
         // Create a new model for high-scoring residues and apply red sticks style
-        let highScoreModel = viewer.addModel(pdb, "pdb");
-        highScoreModel.setStyle({}, {});
-        highScoreModel.setStyle(
+        let class1Model = viewer.addModel(pdb, "pdb");
+        class1Model.setStyle({}, {});
+        class1Model.setStyle(
             {"chain": "%s", "resi": [%s]}, 
-            {"stick": {"color": "red"}}
+            {"stick": {"color": "blue"}}
         );
 
-        // Create a new model for medium-scoring residues and apply orange sticks style
-        let midScoreModel = viewer.addModel(pdb, "pdb");
-        midScoreModel.setStyle({}, {});
-        midScoreModel.setStyle(
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class2Model = viewer.addModel(pdb, "pdb");
+        class2Model.setStyle({}, {});
+        class2Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "lightblue"}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class3Model = viewer.addModel(pdb, "pdb");
+        class3Model.setStyle({}, {});
+        class3Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "white"}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class4Model = viewer.addModel(pdb, "pdb");
+        class4Model.setStyle({}, {});
+        class4Model.setStyle(
             {"chain": "%s", "resi": [%s]}, 
             {"stick": {"color": "orange"}}
         );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class5Model = viewer.addModel(pdb, "pdb");
+        class5Model.setStyle({}, {});
+        class5Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "red"}}
+        );
+
         """ % (
             segment,
             segment,
-            ", ".join(str(resi) for resi in high_score_residues),
+            ", ".join(str(resi) for resi in class1_score_residues),
             segment,
-            ", ".join(str(resi) for resi in mid_score_residues)
+            ", ".join(str(resi) for resi in class2_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class3_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class4_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class5_score_residues)
         )
     
     # Generate the full HTML content
@@ -351,14 +430,22 @@ def molecule(input_pdb, residue_scores=None, segment='A'):
     # Return the HTML content within an iframe safely encoded for special characters
     return f'<iframe width="100%" height="700" srcdoc="{html_content.replace(chr(34), "&quot;").replace(chr(39), "&#39;")}"></iframe>'
 
-
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# Protein Binding Site Prediction")
     
-    with gr.Row():
-        pdb_input = gr.Textbox(value="4BDU", label="PDB ID", placeholder="Enter PDB ID here...")
-        visualize_btn = gr.Button("Visualize Structure")
+    # Mode selection
+    mode = gr.Radio(
+        choices=["PDB ID", "Upload File"],
+        value="PDB ID",
+        label="Input Mode",
+        info="Choose whether to input a PDB ID or upload a PDB/CIF file."
+    )
+
+    # Input components based on mode
+    pdb_input = gr.Textbox(value="4BDU", label="PDB ID", placeholder="Enter PDB ID here...")
+    pdb_file = gr.File(label="Upload PDB/CIF File", visible=False)
+    visualize_btn = gr.Button("Visualize Structure")
 
     molecule_output2 = Molecule3D(label="Protein Structure", reps=[
         {
@@ -375,23 +462,70 @@ with gr.Blocks() as demo:
         segment_input = gr.Textbox(value="A", label="Chain ID", placeholder="Enter Chain ID here...")
         prediction_btn = gr.Button("Predict Binding Site")
 
-
     molecule_output = gr.HTML(label="Protein Structure")
-    predictions_output = gr.Textbox(label="Binding Site Predictions")
+    explanation_vis = gr.Markdown("""
+    Residues with a score > 0.5 are considered binding sites and represented as sticks with the score dependent colorcoding:
+    - 0.5-0.6: blue  
+    - 0.6–0.7: light blue  
+    - 0.7–0.8: white
+    - 0.8–0.9: orange
+    - 0.9–1.0: red
+    """)
+    predictions_output = gr.Textbox(label="Visualize Prediction with PyMol")
+    gr.Markdown("### Download:\n- List of predicted binding site residues\n- PDB with score in beta factor column")
     download_output = gr.File(label="Download Files", file_count="multiple")
     
+    def process_interface(mode, pdb_id, pdb_file, chain_id):
+        if mode == "PDB ID":
+            return process_pdb(pdb_id, chain_id)
+        elif mode == "Upload File":
+            _, ext = os.path.splitext(pdb_file.name)
+            file_path = os.path.join('./', f"{_}{ext}")
+            if ext == '.cif':
+                pdb_path = convert_cif_to_pdb(file_path)
+            else:
+                pdb_path= file_path
+            return process_pdb(pdb_path, chain_id)
+        else:
+            return "Error: Invalid mode selected", None, None
+
+    def fetch_interface(mode, pdb_id, pdb_file):
+        if mode == "PDB ID":
+            return fetch_pdb(pdb_id)
+        elif mode == "Upload File":
+            _, ext = os.path.splitext(pdb_file.name)
+            file_path = os.path.join('./', f"{_}{ext}")
+            #print(ext)
+            if ext == '.cif':
+                pdb_path = convert_cif_to_pdb(file_path)
+            else:
+                pdb_path= file_path
+            #print(pdb_path)
+            return pdb_path
+        else:
+            return "Error: Invalid mode selected"
+
+    def toggle_mode(selected_mode):
+        if selected_mode == "PDB ID":
+            return gr.update(visible=True), gr.update(visible=False)
+        else:
+            return gr.update(visible=False), gr.update(visible=True)
+
+    mode.change(
+        toggle_mode,
+        inputs=[mode],
+        outputs=[pdb_input, pdb_file]
+    )
+
     prediction_btn.click(
-        process_pdb, 
-        inputs=[
-            pdb_input, 
-            segment_input
-        ], 
+        process_interface, 
+        inputs=[mode, pdb_input, pdb_file, segment_input], 
         outputs=[predictions_output, molecule_output, download_output]
     )
 
     visualize_btn.click(
-        fetch_pdb, 
-        inputs=[pdb_input], 
+        fetch_interface, 
+        inputs=[mode, pdb_input, pdb_file], 
         outputs=molecule_output2
     )
 

test.ipynb

@@ -1,846 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "id": "1f8ea359-674c-4263-9c2a-7a8e7e464249",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "* Running on local URL:  http://127.0.0.1:7862\n",
-      "\n",
-      "To create a public link, set `share=True` in `launch()`.\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<div><iframe src=\"http://127.0.0.1:7862/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/plain": []
-     },
-     "execution_count": 3,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import gradio as gr\n",
-    "import requests\n",
-    "from Bio.PDB import PDBParser\n",
-    "from gradio_molecule3d import Molecule3D\n",
-    "import numpy as np\n",
-    "\n",
-    "# Function to fetch a PDB file from RCSB PDB\n",
-    "def fetch_pdb(pdb_id):\n",
-    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
-    "    pdb_path = f'{pdb_id}.pdb'\n",
-    "    response = requests.get(pdb_url)\n",
-    "    if response.status_code == 200:\n",
-    "        with open(pdb_path, 'wb') as f:\n",
-    "            f.write(response.content)\n",
-    "        return pdb_path\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "# Function to process the PDB file and return random predictions\n",
-    "def process_pdb(pdb_id, segment):\n",
-    "    pdb_path = fetch_pdb(pdb_id)\n",
-    "    if not pdb_path:\n",
-    "        return \"Failed to fetch PDB file\", None, None\n",
-    "\n",
-    "    parser = PDBParser(QUIET=True)\n",
-    "    structure = parser.get_structure('protein', pdb_path)\n",
-    "    \n",
-    "    try:\n",
-    "        chain = structure[0][segment]\n",
-    "    except KeyError:\n",
-    "        return \"Invalid Chain ID\", None, None\n",
-    "\n",
-    "    sequence = [residue.get_resname() for residue in chain if residue.id[0] == ' ']\n",
-    "    random_scores = np.random.rand(len(sequence))\n",
-    "\n",
-    "    result_str = \"\\n\".join(\n",
-    "        f\"{seq} {res.id[1]} {score:.2f}\" \n",
-    "        for seq, res, score in zip(sequence, chain, random_scores)\n",
-    "    )\n",
-    "\n",
-    "    # Save the predictions to a file\n",
-    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
-    "    with open(prediction_file, \"w\") as f:\n",
-    "        f.write(result_str)\n",
-    "    \n",
-    "    return result_str, pdb_path, prediction_file\n",
-    "\n",
-    "#reps = [{\"model\": 0, \"style\": \"cartoon\", \"color\": \"spectrum\"}]\n",
-    "\n",
-    "reps =    [\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"style\": \"cartoon\",\n",
-    "          \"color\": \"whiteCarbon\",\n",
-    "          \"residue_range\": \"\",\n",
-    "          \"around\": 0,\n",
-    "          \"byres\": False,\n",
-    "        },\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"chain\": \"A\",\n",
-    "          \"resname\": \"HIS\",\n",
-    "          \"style\": \"stick\",\n",
-    "          \"color\": \"red\"\n",
-    "        }\n",
-    "      ]\n",
-    "\n",
-    "\n",
-    "# Gradio UI\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
-    "\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
-    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
-    "\n",
-    "    molecule_output = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
-    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
-    "    download_output = gr.File(label=\"Download Predictions\")\n",
-    "\n",
-    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output)\n",
-    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
-    "\n",
-    "    gr.Markdown(\"## Examples\")\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            [\"2IWI\", \"A\"],\n",
-    "            [\"7RPZ\", \"B\"],\n",
-    "            [\"3TJN\", \"C\"]\n",
-    "        ],\n",
-    "        inputs=[pdb_input, segment_input],\n",
-    "        outputs=[predictions_output, molecule_output, download_output]\n",
-    "    )\n",
-    "\n",
-    "demo.launch()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "bd50ff2e-ed03-498e-8af2-73c0fb8ea07e",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "raw",
-   "id": "88affe12-7c48-4bd6-9e46-32cdffa729fe",
-   "metadata": {},
-   "source": [
-    "import gradio as gr\n",
-    "from gradio_molecule3d import Molecule3D\n",
-    "\n",
-    "\n",
-    "example = Molecule3D().example_value()\n",
-    "\n",
-    "\n",
-    "reps =    [\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"style\": \"cartoon\",\n",
-    "          \"color\": \"whiteCarbon\",\n",
-    "          \"residue_range\": \"\",\n",
-    "          \"around\": 0,\n",
-    "          \"byres\": False,\n",
-    "        },\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"chain\": \"A\",\n",
-    "          \"resname\": \"HIS\",\n",
-    "          \"style\": \"stick\",\n",
-    "          \"color\": \"red\"\n",
-    "        }\n",
-    "      ]\n",
-    "\n",
-    "\n",
-    "\n",
-    "def predict(x):\n",
-    "    print(\"predict function\", x)\n",
-    "    print(x.name)\n",
-    "    return x\n",
-    "\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Molecule3D\")\n",
-    "    inp = Molecule3D(label=\"Molecule3D\", reps=reps)\n",
-    "    out = Molecule3D(label=\"Output\", reps=reps)\n",
-    "\n",
-    "    btn = gr.Button(\"Predict\")\n",
-    "    gr.Markdown(\"\"\" \n",
-    "    You can configure the default rendering of the molecule by adding a list of representations\n",
-    "    <pre>\n",
-    "        reps =    [\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"style\": \"cartoon\",\n",
-    "          \"color\": \"whiteCarbon\",\n",
-    "          \"residue_range\": \"\",\n",
-    "          \"around\": 0,\n",
-    "          \"byres\": False,\n",
-    "        },\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"chain\": \"A\",\n",
-    "          \"resname\": \"HIS\",\n",
-    "          \"style\": \"stick\",\n",
-    "          \"color\": \"red\"\n",
-    "        }\n",
-    "      ]\n",
-    "    </pre>\n",
-    "    \"\"\")\n",
-    "    btn.click(predict, inputs=inp, outputs=out)\n",
-    "\n",
-    "\n",
-    "if __name__ == \"__main__\":\n",
-    "    demo.launch()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "d27cc368-26a0-42c2-a68a-8833de7bb4a0",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "raw",
-   "id": "2b970adb-3152-427f-bb58-b92974ff406e",
-   "metadata": {},
-   "source": [
-    "import gradio as gr\n",
-    "import os\n",
-    "import requests\n",
-    "from Bio.PDB import PDBParser, PDBIO\n",
-    "import biotite.structure.io as bsio\n",
-    "\n",
-    "def read_mol(pdb_path):\n",
-    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
-    "    with open(pdb_path, 'r') as f:\n",
-    "        return f.read()\n",
-    "\n",
-    "# Function to fetch or upload the PDB file\n",
-    "def get_pdb(pdb_code=\"\", filepath=\"\"):\n",
-    "    if pdb_code and len(pdb_code) == 4:\n",
-    "        pdb_file = f\"{pdb_code}.pdb\"\n",
-    "        if not os.path.exists(pdb_file):\n",
-    "            os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
-    "        return pdb_file\n",
-    "    elif filepath is not None:\n",
-    "        return filepath\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "def molecule(input_pdb):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
-    "    \n",
-    "    html_content = f\"\"\"\n",
-    "    <!DOCTYPE html>\n",
-    "    <html>\n",
-    "    <head>    \n",
-    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
-    "        <style>\n",
-    "        .mol-container {{\n",
-    "            width: 100%;\n",
-    "            height: 700px;\n",
-    "            position: relative;\n",
-    "        }}\n",
-    "        </style>\n",
-    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
-    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
-    "    </head>\n",
-    "    <body>\n",
-    "        <div id=\"container\" class=\"mol-container\"></div>\n",
-    "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
-    "            $(document).ready(function () {{\n",
-    "                let element = $(\"#container\");\n",
-    "                let config = {{ backgroundColor: \"white\" }};\n",
-    "                let viewer = $3Dmol.createViewer(element, config);\n",
-    "                viewer.addModel(pdb, \"pdb\");\n",
-    "                viewer.getModel(0).setStyle({{}}, {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }});\n",
-    "                viewer.zoomTo();\n",
-    "                viewer.render();\n",
-    "                viewer.zoom(0.8, 2000);\n",
-    "            }});\n",
-    "        </script>\n",
-    "    </body>\n",
-    "    </html>\n",
-    "    \"\"\"\n",
-    "    \n",
-    "    # Return the HTML content within an iframe safely encoded for special characters\n",
-    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
-    "\n",
-    "# Gradio function to update the visualization\n",
-    "def update(inp, file):\n",
-    "    pdb_path = get_pdb(inp, file)\n",
-    "    if pdb_path:\n",
-    "        return molecule(pdb_path)\n",
-    "    else:\n",
-    "        return \"Invalid input. Please provide a valid PDB code or upload a PDB file.\"\n",
-    "\n",
-    "# Gradio UI\n",
-    "demo = gr.Blocks()\n",
-    "with demo:\n",
-    "    gr.Markdown(\"# PDB Viewer using 3Dmol.js\")\n",
-    "    with gr.Row():\n",
-    "        with gr.Column():\n",
-    "            inp = gr.Textbox(\n",
-    "                placeholder=\"PDB Code or upload file below\", label=\"Input structure\"\n",
-    "            )\n",
-    "            file = gr.File(file_count=\"single\")\n",
-    "            btn = gr.Button(\"View structure\")\n",
-    "        mol = gr.HTML()\n",
-    "    btn.click(fn=update, inputs=[inp, file], outputs=mol)\n",
-    "\n",
-    "# Launch the Gradio interface \n",
-    "demo.launch(debug=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "ee215c16-a1fb-450f-bb93-37aaee6fb3f1",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "raw",
-   "id": "050aa2e8-2dbe-4a28-8692-58ca7c50fccd",
-   "metadata": {},
-   "source": [
-    "import gradio as gr\n",
-    "import os\n",
-    "import requests\n",
-    "import numpy as np\n",
-    "from Bio.PDB import PDBParser\n",
-    "\n",
-    "def read_mol(pdb_path):\n",
-    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
-    "    with open(pdb_path, 'r') as f:\n",
-    "        return f.read()\n",
-    "\n",
-    "# Function to fetch a PDB file from RCSB PDB\n",
-    "def fetch_pdb(pdb_id):\n",
-    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
-    "    pdb_path = f'{pdb_id}.pdb'\n",
-    "    response = requests.get(pdb_url)\n",
-    "    if response.status_code == 200:\n",
-    "        with open(pdb_path, 'wb') as f:\n",
-    "            f.write(response.content)\n",
-    "        return molecule(pdb_path)\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "# Function to process the PDB file and return random predictions\n",
-    "def process_pdb(pdb_id, segment):\n",
-    "    pdb_path = fetch_pdb(pdb_id)\n",
-    "    if not pdb_path:\n",
-    "        return \"Failed to fetch PDB file\", None, None\n",
-    "    \n",
-    "    parser = PDBParser(QUIET=True)\n",
-    "    structure = parser.get_structure('protein', pdb_path)\n",
-    "    \n",
-    "    try:\n",
-    "        chain = structure[0][segment]\n",
-    "    except KeyError:\n",
-    "        return \"Invalid Chain ID\", None, None\n",
-    "    \n",
-    "    sequence = [residue.get_resname() for residue in chain if residue.id[0] == ' ']\n",
-    "    random_scores = np.random.rand(len(sequence))\n",
-    "    result_str = \"\\n\".join(\n",
-    "        f\"{seq} {res.id[1]} {score:.2f}\" \n",
-    "        for seq, res, score in zip(sequence, chain, random_scores)\n",
-    "    )\n",
-    "    \n",
-    "    # Save the predictions to a file\n",
-    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
-    "    with open(prediction_file, \"w\") as f:\n",
-    "        f.write(result_str)\n",
-    "    \n",
-    "    return result_str, molecule(pdb_path), prediction_file\n",
-    "\n",
-    "def molecule(input_pdb):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
-    "    \n",
-    "    html_content = f\"\"\"\n",
-    "    <!DOCTYPE html>\n",
-    "    <html>\n",
-    "    <head>    \n",
-    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
-    "        <style>\n",
-    "        .mol-container {{\n",
-    "            width: 100%;\n",
-    "            height: 700px;\n",
-    "            position: relative;\n",
-    "        }}\n",
-    "        </style>\n",
-    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
-    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
-    "    </head>\n",
-    "    <body>\n",
-    "        <div id=\"container\" class=\"mol-container\"></div>\n",
-    "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
-    "            $(document).ready(function () {{\n",
-    "                let element = $(\"#container\");\n",
-    "                let config = {{ backgroundColor: \"white\" }};\n",
-    "                let viewer = $3Dmol.createViewer(element, config);\n",
-    "                viewer.addModel(pdb, \"pdb\");\n",
-    "                \n",
-    "                // Set cartoon representation with white carbon color scheme\n",
-    "                viewer.getModel(0).setStyle({{}}, {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }});\n",
-    "                \n",
-    "                // Highlight specific histidine residues in red stick representation\n",
-    "                viewer.getModel(0).setStyle(\n",
-    "                    {{\"resn\": \"HIS\"}}, \n",
-    "                    {{\"stick\": {{\"color\": \"red\"}}}}\n",
-    "                );\n",
-    "                \n",
-    "                viewer.zoomTo();\n",
-    "                viewer.render();\n",
-    "                viewer.zoom(0.8, 2000);\n",
-    "            }});\n",
-    "        </script>\n",
-    "    </body>\n",
-    "    </html>\n",
-    "    \"\"\"\n",
-    "    \n",
-    "    # Return the HTML content within an iframe safely encoded for special characters\n",
-    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
-    "\n",
-    "# Gradio UI\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
-    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
-    "    \n",
-    "    # Use HTML output instead of Molecule3D\n",
-    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
-    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
-    "    download_output = gr.File(label=\"Download Predictions\")\n",
-    "    \n",
-    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output)\n",
-    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
-    "    \n",
-    "    gr.Markdown(\"## Examples\")\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            [\"2IWI\", \"A\"],\n",
-    "            [\"7RPZ\", \"B\"],\n",
-    "            [\"3TJN\", \"C\"]\n",
-    "        ],\n",
-    "        inputs=[pdb_input, segment_input],\n",
-    "        outputs=[predictions_output, molecule_output, download_output]\n",
-    "    )\n",
-    "\n",
-    "demo.launch(debug=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "9a5facd9-855c-4b35-8dd3-2c0c8c7dd356",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "raw",
-   "id": "a762170f-92a9-473d-b18d-53607a780e3b",
-   "metadata": {},
-   "source": [
-    "import gradio as gr\n",
-    "import requests\n",
-    "from Bio.PDB import PDBParser\n",
-    "import numpy as np\n",
-    "import os\n",
-    "\n",
-    "def read_mol(pdb_path):\n",
-    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
-    "    with open(pdb_path, 'r') as f:\n",
-    "        return f.read()\n",
-    "\n",
-    "# Function to fetch a PDB file from RCSB PDB\n",
-    "def fetch_pdb(pdb_id):\n",
-    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
-    "    pdb_path = f'{pdb_id}.pdb'\n",
-    "    response = requests.get(pdb_url)\n",
-    "    if response.status_code == 200:\n",
-    "        with open(pdb_path, 'wb') as f:\n",
-    "            f.write(response.content)\n",
-    "        return pdb_path\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "# Function to process the PDB file and return random predictions\n",
-    "def process_pdb(pdb_id, segment):\n",
-    "    pdb_path = fetch_pdb(pdb_id)\n",
-    "    if not pdb_path:\n",
-    "        return \"Failed to fetch PDB file\", None, None\n",
-    "    parser = PDBParser(QUIET=True)\n",
-    "    structure = parser.get_structure('protein', pdb_path)\n",
-    "    \n",
-    "    try:\n",
-    "        chain = structure[0][segment]\n",
-    "    except KeyError:\n",
-    "        return \"Invalid Chain ID\", None, None\n",
-    "    sequence = [residue.get_resname() for residue in chain if residue.id[0] == ' ']\n",
-    "    random_scores = np.random.rand(len(sequence))\n",
-    "    result_str = \"\\n\".join(\n",
-    "        f\"{seq} {res.id[1]} {score:.2f}\" \n",
-    "        for seq, res, score in zip(sequence, chain, random_scores)\n",
-    "    )\n",
-    "    # Save the predictions to a file\n",
-    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
-    "    with open(prediction_file, \"w\") as f:\n",
-    "        f.write(result_str)\n",
-    "    \n",
-    "    return result_str, molecule(pdb_path), prediction_file\n",
-    "\n",
-    "def molecule(input_pdb):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
-    "    \n",
-    "    html_content = f\"\"\"\n",
-    "    <!DOCTYPE html>\n",
-    "    <html>\n",
-    "    <head>    \n",
-    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
-    "        <style>\n",
-    "        .mol-container {{\n",
-    "            width: 100%;\n",
-    "            height: 700px;\n",
-    "            position: relative;\n",
-    "        }}\n",
-    "        </style>\n",
-    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
-    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
-    "    </head>\n",
-    "    <body>\n",
-    "        <div id=\"container\" class=\"mol-container\"></div>\n",
-    "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
-    "            $(document).ready(function () {{\n",
-    "                let element = $(\"#container\");\n",
-    "                let config = {{ backgroundColor: \"white\" }};\n",
-    "                let viewer = $3Dmol.createViewer(element, config);\n",
-    "                viewer.addModel(pdb, \"pdb\");\n",
-    "                \n",
-    "                // Set cartoon representation with white carbon color scheme\n",
-    "                viewer.getModel(0).setStyle({{}}, {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }});\n",
-    "                \n",
-    "                // Highlight specific histidine residues in red stick representation\n",
-    "                viewer.getModel(0).setStyle(\n",
-    "                    {{\"resn\": \"HIS\"}}, \n",
-    "                    {{\"stick\": {{\"color\": \"red\"}}}}\n",
-    "                );\n",
-    "                \n",
-    "                viewer.zoomTo();\n",
-    "                viewer.render();\n",
-    "                viewer.zoom(0.8, 2000);\n",
-    "            }});\n",
-    "        </script>\n",
-    "    </body>\n",
-    "    </html>\n",
-    "    \"\"\"\n",
-    "    \n",
-    "    # Return the HTML content within an iframe safely encoded for special characters\n",
-    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
-    "\n",
-    "# Gradio UI\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
-    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
-    "    \n",
-    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
-    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
-    "    download_output = gr.File(label=\"Download Predictions\")\n",
-    "    \n",
-    "    # Update to explicitly use molecule() function for visualization\n",
-    "    visualize_btn.click(\n",
-    "        fn=lambda pdb_id: molecule(fetch_pdb(pdb_id)), \n",
-    "        inputs=[pdb_input], \n",
-    "        outputs=molecule_output\n",
-    "    )\n",
-    "    \n",
-    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
-    "    \n",
-    "    gr.Markdown(\"## Examples\")\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            [\"2IWI\", \"A\"],\n",
-    "            [\"7RPZ\", \"B\"],\n",
-    "            [\"3TJN\", \"C\"]\n",
-    "        ],\n",
-    "        inputs=[pdb_input, segment_input],\n",
-    "        outputs=[predictions_output, molecule_output, download_output]\n",
-    "    )\n",
-    "\n",
-    "demo.launch()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "15527a58-c449-4da0-8fab-3baaede15e41",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "9ef3e330-cb88-4c29-b84a-2f8652883cfc",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "* Running on local URL:  http://127.0.0.1:7860\n",
-      "\n",
-      "To create a public link, set `share=True` in `launch()`.\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<div><iframe src=\"http://127.0.0.1:7860/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/plain": []
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import gradio as gr\n",
-    "import requests\n",
-    "from Bio.PDB import PDBParser\n",
-    "import numpy as np\n",
-    "import os\n",
-    "from gradio_molecule3d import Molecule3D\n",
-    "\n",
-    "def read_mol(pdb_path):\n",
-    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
-    "    with open(pdb_path, 'r') as f:\n",
-    "        return f.read()\n",
-    "\n",
-    "def fetch_pdb(pdb_id):\n",
-    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
-    "    pdb_path = f'{pdb_id}.pdb'\n",
-    "    response = requests.get(pdb_url)\n",
-    "    if response.status_code == 200:\n",
-    "        with open(pdb_path, 'wb') as f:\n",
-    "            f.write(response.content)\n",
-    "        return pdb_path\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "def process_pdb(pdb_id, segment):\n",
-    "    pdb_path = fetch_pdb(pdb_id)\n",
-    "    if not pdb_path:\n",
-    "        return \"Failed to fetch PDB file\", None, None\n",
-    "    parser = PDBParser(QUIET=True)\n",
-    "    structure = parser.get_structure('protein', pdb_path)\n",
-    "    \n",
-    "    try:\n",
-    "        chain = structure[0][segment]\n",
-    "    except KeyError:\n",
-    "        return \"Invalid Chain ID\", None, None\n",
-    "    sequence = [residue.get_resname() for residue in chain if residue.id[0] == ' ']\n",
-    "    random_scores = np.random.rand(len(sequence))\n",
-    "    result_str = \"\\n\".join(\n",
-    "        f\"{seq} {res.id[1]} {score:.2f}\" \n",
-    "        for seq, res, score in zip(sequence, chain, random_scores)\n",
-    "    )\n",
-    "    # Save the predictions to a file\n",
-    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
-    "    with open(prediction_file, \"w\") as f:\n",
-    "        f.write(result_str)\n",
-    "    \n",
-    "    return result_str, molecule(pdb_path, random_scores), prediction_file\n",
-    "\n",
-    "def molecule(input_pdb, scores=None):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
-    "    \n",
-    "    # Prepare high-scoring residues script if scores are provided\n",
-    "    high_score_script = \"\"\n",
-    "    if scores is not None:\n",
-    "        high_score_script = \"\"\"\n",
-    "        // Highlight residues with high scores\n",
-    "        let highScoreResidues = [{}];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {{\"resi\": highScoreResidues}}, \n",
-    "            {{\"stick\": {{\"color\": \"red\"}}}}\n",
-    "        );\n",
-    "        \"\"\".format(\n",
-    "            \", \".join(str(i+1) for i, score in enumerate(scores) if score > 0.8)\n",
-    "        )\n",
-    "    \n",
-    "    html_content = f\"\"\"\n",
-    "    <!DOCTYPE html>\n",
-    "    <html>\n",
-    "    <head>    \n",
-    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
-    "        <style>\n",
-    "        .mol-container {{\n",
-    "            width: 100%;\n",
-    "            height: 700px;\n",
-    "            position: relative;\n",
-    "        }}\n",
-    "        </style>\n",
-    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
-    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
-    "    </head>\n",
-    "    <body>\n",
-    "        <div id=\"container\" class=\"mol-container\"></div>\n",
-    "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
-    "            $(document).ready(function () {{\n",
-    "                let element = $(\"#container\");\n",
-    "                let config = {{ backgroundColor: \"white\" }};\n",
-    "                let viewer = $3Dmol.createViewer(element, config);\n",
-    "                viewer.addModel(pdb, \"pdb\");\n",
-    "                \n",
-    "                // Set cartoon representation with white carbon color scheme\n",
-    "                viewer.getModel(0).setStyle({{}}, {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }});\n",
-    "                \n",
-    "                {high_score_script}\n",
-    "                \n",
-    "                viewer.zoomTo();\n",
-    "                viewer.render();\n",
-    "                viewer.zoom(0.8, 2000);\n",
-    "            }});\n",
-    "        </script>\n",
-    "    </body>\n",
-    "    </html>\n",
-    "    \"\"\"\n",
-    "    \n",
-    "    # Return the HTML content within an iframe safely encoded for special characters\n",
-    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
-    "\n",
-    "reps =    [\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"style\": \"cartoon\",\n",
-    "          \"color\": \"whiteCarbon\",\n",
-    "          \"residue_range\": \"\",\n",
-    "          \"around\": 0,\n",
-    "          \"byres\": False,\n",
-    "        }\n",
-    "    ]\n",
-    "# Gradio UI\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
-    "        #prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
-    "\n",
-    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
-    "\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
-    "\n",
-    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
-    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
-    "    download_output = gr.File(label=\"Download Predictions\")\n",
-    "    \n",
-    "    #visualize_btn.click(\n",
-    "    #    fn=lambda pdb_id: molecule(fetch_pdb(pdb_id)), \n",
-    "    #    inputs=[pdb_input], \n",
-    "    #    outputs=molecule_output\n",
-    "    #)\n",
-    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)\n",
-    "    \n",
-    "    \n",
-    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
-    "    \n",
-    "    gr.Markdown(\"## Examples\")\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            [\"2IWI\", \"A\"],\n",
-    "            [\"7RPZ\", \"B\"],\n",
-    "            [\"3TJN\", \"C\"]\n",
-    "        ],\n",
-    "        inputs=[pdb_input, segment_input],\n",
-    "        outputs=[predictions_output, molecule_output, download_output]\n",
-    "    )\n",
-    "\n",
-    "demo.launch()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "14605615-8610-4d9e-841b-db7618cde844",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python (LLM)",
-   "language": "python",
-   "name": "llm"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.12.7"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

test2.ipynb

@@ -1,1598 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "f3b7f6b0-6685-4a5c-9529-45e0ca905a3b",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "* Running on local URL:  http://127.0.0.1:7860\n",
-      "\n",
-      "To create a public link, set `share=True` in `launch()`.\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<div><iframe src=\"http://127.0.0.1:7860/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/plain": []
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import gradio as gr\n",
-    "import requests\n",
-    "from Bio.PDB import PDBParser\n",
-    "import numpy as np\n",
-    "import os\n",
-    "from gradio_molecule3d import Molecule3D\n",
-    "\n",
-    "def read_mol(pdb_path):\n",
-    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
-    "    with open(pdb_path, 'r') as f:\n",
-    "        return f.read()\n",
-    "\n",
-    "def fetch_pdb(pdb_id):\n",
-    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
-    "    pdb_path = f'{pdb_id}.pdb'\n",
-    "    response = requests.get(pdb_url)\n",
-    "    if response.status_code == 200:\n",
-    "        with open(pdb_path, 'wb') as f:\n",
-    "            f.write(response.content)\n",
-    "        return pdb_path\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "def process_pdb(pdb_id, segment):\n",
-    "    pdb_path = fetch_pdb(pdb_id)\n",
-    "    if not pdb_path:\n",
-    "        return \"Failed to fetch PDB file\", None, None\n",
-    "    \n",
-    "    parser = PDBParser(QUIET=1)\n",
-    "    structure = parser.get_structure('protein', pdb_path)\n",
-    "    \n",
-    "    try:\n",
-    "        chain = structure[0][segment]\n",
-    "    except KeyError:\n",
-    "        return \"Invalid Chain ID\", None, None\n",
-    "    \n",
-    "    # Comprehensive amino acid mapping\n",
-    "    aa_dict = {\n",
-    "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
-    "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
-    "        'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',\n",
-    "        'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',\n",
-    "        'MSE': 'M', 'SEP': 'S', 'TPO': 'T', 'CSO': 'C', 'PTR': 'Y', 'HYP': 'P'\n",
-    "    }\n",
-    "    \n",
-    "    # Exclude non-amino acid residues\n",
-    "    sequence = [\n",
-    "        residue for residue in chain \n",
-    "        if residue.get_resname().strip() in aa_dict\n",
-    "    ]\n",
-    "    \n",
-    "    random_scores = np.random.rand(len(sequence))\n",
-    "    result_str = \"\\n\".join(\n",
-    "        f\"{aa_dict[res.get_resname()]} {res.id[1]} {score:.2f}\" \n",
-    "        for res, score in zip(sequence, random_scores)\n",
-    "    )\n",
-    "    \n",
-    "    # Save the predictions to a file\n",
-    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
-    "    with open(prediction_file, \"w\") as f:\n",
-    "        f.write(result_str)\n",
-    "    \n",
-    "    return result_str, molecule(pdb_path, random_scores, segment), prediction_file\n",
-    "\n",
-    "def molecule(input_pdb, scores=None, segment='A'):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
-    "    \n",
-    "    # Prepare high-scoring residues script if scores are provided\n",
-    "    high_score_script = \"\"\n",
-    "    if scores is not None:\n",
-    "        high_score_script = \"\"\"\n",
-    "        // Reset all styles first\n",
-    "        viewer.getModel(0).setStyle({}, {});\n",
-    "        \n",
-    "        // Show only the selected chain\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\"}, \n",
-    "            { cartoon: {colorscheme:\"whiteCarbon\"} }\n",
-    "        );\n",
-    "        \n",
-    "        // Highlight high-scoring residues only for the selected chain\n",
-    "        let highScoreResidues = [%s];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": highScoreResidues}, \n",
-    "            {\"stick\": {\"color\": \"red\"}}\n",
-    "        );\n",
-    "        \"\"\" % (segment, \n",
-    "               \", \".join(str(i+1) for i, score in enumerate(scores) if score > 0.8),\n",
-    "               segment)\n",
-    "    \n",
-    "    html_content = f\"\"\"\n",
-    "    <!DOCTYPE html>\n",
-    "    <html>\n",
-    "    <head>    \n",
-    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
-    "        <style>\n",
-    "        .mol-container {{\n",
-    "            width: 100%;\n",
-    "            height: 700px;\n",
-    "            position: relative;\n",
-    "        }}\n",
-    "        </style>\n",
-    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
-    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
-    "    </head>\n",
-    "    <body>\n",
-    "        <div id=\"container\" class=\"mol-container\"></div>\n",
-    "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
-    "            $(document).ready(function () {{\n",
-    "                let element = $(\"#container\");\n",
-    "                let config = {{ backgroundColor: \"white\" }};\n",
-    "                let viewer = $3Dmol.createViewer(element, config);\n",
-    "                viewer.addModel(pdb, \"pdb\");\n",
-    "                \n",
-    "                // Reset all styles and show only selected chain\n",
-    "                viewer.getModel(0).setStyle(\n",
-    "                    {{\"chain\": \"{segment}\"}}, \n",
-    "                    {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }}\n",
-    "                );\n",
-    "                \n",
-    "                {high_score_script}\n",
-    "                \n",
-    "                viewer.zoomTo();\n",
-    "                viewer.render();\n",
-    "                viewer.zoom(0.8, 2000);\n",
-    "            }});\n",
-    "        </script>\n",
-    "    </body>\n",
-    "    </html>\n",
-    "    \"\"\"\n",
-    "    \n",
-    "    # Return the HTML content within an iframe safely encoded for special characters\n",
-    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
-    "\n",
-    "reps =    [\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"style\": \"cartoon\",\n",
-    "          \"color\": \"whiteCarbon\",\n",
-    "          \"residue_range\": \"\",\n",
-    "          \"around\": 0,\n",
-    "          \"byres\": False,\n",
-    "        }\n",
-    "    ]\n",
-    "# Gradio UI\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
-    "\n",
-    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
-    "\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
-    "\n",
-    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
-    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
-    "    download_output = gr.File(label=\"Download Predictions\")\n",
-    "    \n",
-    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)\n",
-    "    \n",
-    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
-    "    \n",
-    "    gr.Markdown(\"## Examples\")\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            [\"2IWI\", \"A\"],\n",
-    "            [\"7RPZ\", \"B\"],\n",
-    "            [\"3TJN\", \"C\"]\n",
-    "        ],\n",
-    "        inputs=[pdb_input, segment_input],\n",
-    "        outputs=[predictions_output, molecule_output, download_output]\n",
-    "    )\n",
-    "\n",
-    "demo.launch()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "28f8f28c-48d3-4e35-9766-3de9882179b5",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "* Running on local URL:  http://127.0.0.1:7864\n",
-      "\n",
-      "To create a public link, set `share=True` in `launch()`.\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<div><iframe src=\"http://127.0.0.1:7864/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/plain": []
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import gradio as gr\n",
-    "import requests\n",
-    "from Bio.PDB import PDBParser\n",
-    "import numpy as np\n",
-    "import os\n",
-    "from gradio_molecule3d import Molecule3D\n",
-    "\n",
-    "def read_mol(pdb_path):\n",
-    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
-    "    with open(pdb_path, 'r') as f:\n",
-    "        return f.read()\n",
-    "\n",
-    "def fetch_pdb(pdb_id):\n",
-    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
-    "    pdb_path = f'{pdb_id}.pdb'\n",
-    "    response = requests.get(pdb_url)\n",
-    "    if response.status_code == 200:\n",
-    "        with open(pdb_path, 'wb') as f:\n",
-    "            f.write(response.content)\n",
-    "        return pdb_path\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "def process_pdb(pdb_id, segment):\n",
-    "    pdb_path = fetch_pdb(pdb_id)\n",
-    "    if not pdb_path:\n",
-    "        return \"Failed to fetch PDB file\", None, None\n",
-    "    \n",
-    "    parser = PDBParser(QUIET=1)\n",
-    "    structure = parser.get_structure('protein', pdb_path)\n",
-    "    \n",
-    "    try:\n",
-    "        chain = structure[0][segment]\n",
-    "    except KeyError:\n",
-    "        return \"Invalid Chain ID\", None, None\n",
-    "    \n",
-    "    # Comprehensive amino acid mapping\n",
-    "    aa_dict = {\n",
-    "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
-    "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
-    "        'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',\n",
-    "        'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',\n",
-    "        'MSE': 'M', 'SEP': 'S', 'TPO': 'T', 'CSO': 'C', 'PTR': 'Y', 'HYP': 'P'\n",
-    "    }\n",
-    "    \n",
-    "    # Exclude non-amino acid residues\n",
-    "    sequence = [\n",
-    "        residue for residue in chain \n",
-    "        if residue.get_resname().strip() in aa_dict\n",
-    "    ]\n",
-    "    \n",
-    "    random_scores = np.random.rand(len(sequence))\n",
-    "    result_str = \"\\n\".join(\n",
-    "        f\"{aa_dict[res.get_resname()]} {res.id[1]} {score:.2f}\" \n",
-    "        for res, score in zip(sequence, random_scores)\n",
-    "    )\n",
-    "    \n",
-    "    # Save the predictions to a file\n",
-    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
-    "    with open(prediction_file, \"w\") as f:\n",
-    "        f.write(result_str)\n",
-    "    \n",
-    "    return result_str, molecule(pdb_path, random_scores, segment), prediction_file\n",
-    "\n",
-    "def molecule(input_pdb, scores=None, segment='A'):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
-    "    \n",
-    "    # Prepare high-scoring residues script if scores are provided\n",
-    "    high_score_script = \"\"\n",
-    "    if scores is not None:\n",
-    "        high_score_script = \"\"\"\n",
-    "        // Reset all styles first\n",
-    "        viewer.getModel(0).setStyle({}, {});\n",
-    "        \n",
-    "        // Show only the selected chain\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\"}, \n",
-    "            { cartoon: {colorscheme:\"whiteCarbon\"} }\n",
-    "        );\n",
-    "        \n",
-    "        // Highlight high-scoring residues only for the selected chain\n",
-    "        let highScoreResidues = [%s];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": highScoreResidues}, \n",
-    "            {\"stick\": {\"color\": \"red\"}}\n",
-    "        );\n",
-    "        \"\"\" % (segment, \n",
-    "               \", \".join(str(i+1) for i, score in enumerate(scores) if score > 0.8),\n",
-    "               segment)\n",
-    "    \n",
-    "    html_content = f\"\"\"\n",
-    "    <!DOCTYPE html>\n",
-    "    <html>\n",
-    "    <head>    \n",
-    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
-    "        <style>\n",
-    "        .mol-container {{\n",
-    "            width: 100%;\n",
-    "            height: 700px;\n",
-    "            position: relative;\n",
-    "        }}\n",
-    "        </style>\n",
-    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
-    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
-    "    </head>\n",
-    "    <body>\n",
-    "        <div id=\"container\" class=\"mol-container\"></div>\n",
-    "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
-    "            $(document).ready(function () {{\n",
-    "                let element = $(\"#container\");\n",
-    "                let config = {{ backgroundColor: \"white\" }};\n",
-    "                let viewer = $3Dmol.createViewer(element, config);\n",
-    "                viewer.addModel(pdb, \"pdb\");\n",
-    "                \n",
-    "                // Reset all styles and show only selected chain\n",
-    "                viewer.getModel(0).setStyle(\n",
-    "                    {{\"chain\": \"{segment}\"}}, \n",
-    "                    {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }}\n",
-    "                );\n",
-    "                \n",
-    "                {high_score_script}\n",
-    "                \n",
-    "                // Add hover functionality\n",
-    "                viewer.setHoverable(\n",
-    "                    {{}}, \n",
-    "                    true, \n",
-    "                    function(atom, viewer, event, container) {{\n",
-    "                        if (!atom.label) {{\n",
-    "                            atom.label = viewer.addLabel(\n",
-    "                                atom.resn + \":\" + atom.atom, \n",
-    "                                {{\n",
-    "                                    position: atom, \n",
-    "                                    backgroundColor: 'mintcream', \n",
-    "                                    fontColor: 'black',\n",
-    "                                    fontSize: 12,\n",
-    "                                    padding: 2\n",
-    "                                }}\n",
-    "                            );\n",
-    "                        }}\n",
-    "                    }},\n",
-    "                    function(atom, viewer) {{\n",
-    "                        if (atom.label) {{\n",
-    "                            viewer.removeLabel(atom.label);\n",
-    "                            delete atom.label;\n",
-    "                        }}\n",
-    "                    }}\n",
-    "                );\n",
-    "                \n",
-    "                viewer.zoomTo();\n",
-    "                viewer.render();\n",
-    "                viewer.zoom(0.8, 2000);\n",
-    "            }});\n",
-    "        </script>\n",
-    "    </body>\n",
-    "    </html>\n",
-    "    \"\"\"\n",
-    "    \n",
-    "    # Return the HTML content within an iframe safely encoded for special characters\n",
-    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
-    "\n",
-    "reps =    [\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"style\": \"cartoon\",\n",
-    "          \"color\": \"whiteCarbon\",\n",
-    "          \"residue_range\": \"\",\n",
-    "          \"around\": 0,\n",
-    "          \"byres\": False,\n",
-    "        }\n",
-    "    ]\n",
-    "\n",
-    "# Gradio UI\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
-    "\n",
-    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
-    "\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
-    "\n",
-    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
-    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
-    "    download_output = gr.File(label=\"Download Predictions\")\n",
-    "    \n",
-    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)\n",
-    "    \n",
-    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
-    "    \n",
-    "    gr.Markdown(\"## Examples\")\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            [\"2IWI\", \"A\"],\n",
-    "            [\"7RPZ\", \"B\"],\n",
-    "            [\"3TJN\", \"C\"]\n",
-    "        ],\n",
-    "        inputs=[pdb_input, segment_input],\n",
-    "        outputs=[predictions_output, molecule_output, download_output]\n",
-    "    )\n",
-    "\n",
-    "demo.launch()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "517a2fe7-419f-4d0b-a9ed-62a22c1c1284",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "d62be1b5-762e-4b69-aed4-e4ba2a44482f",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "* Running on local URL:  http://127.0.0.1:7860\n",
-      "\n",
-      "To create a public link, set `share=True` in `launch()`.\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<div><iframe src=\"http://127.0.0.1:7860/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/plain": []
-     },
-     "execution_count": 1,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import gradio as gr\n",
-    "import requests\n",
-    "from Bio.PDB import PDBParser\n",
-    "import numpy as np\n",
-    "import os\n",
-    "from gradio_molecule3d import Molecule3D\n",
-    "\n",
-    "def read_mol(pdb_path):\n",
-    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
-    "    with open(pdb_path, 'r') as f:\n",
-    "        return f.read()\n",
-    "\n",
-    "def fetch_pdb(pdb_id):\n",
-    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
-    "    pdb_path = f'{pdb_id}.pdb'\n",
-    "    response = requests.get(pdb_url)\n",
-    "    if response.status_code == 200:\n",
-    "        with open(pdb_path, 'wb') as f:\n",
-    "            f.write(response.content)\n",
-    "        return pdb_path\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "def process_pdb(pdb_id, segment):\n",
-    "    pdb_path = fetch_pdb(pdb_id)\n",
-    "    if not pdb_path:\n",
-    "        return \"Failed to fetch PDB file\", None, None\n",
-    "    \n",
-    "    parser = PDBParser(QUIET=1)\n",
-    "    structure = parser.get_structure('protein', pdb_path)\n",
-    "    \n",
-    "    try:\n",
-    "        chain = structure[0][segment]\n",
-    "    except KeyError:\n",
-    "        return \"Invalid Chain ID\", None, None\n",
-    "    \n",
-    "    # Comprehensive amino acid mapping\n",
-    "    aa_dict = {\n",
-    "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
-    "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
-    "        'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',\n",
-    "        'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',\n",
-    "        'MSE': 'M', 'SEP': 'S', 'TPO': 'T', 'CSO': 'C', 'PTR': 'Y', 'HYP': 'P'\n",
-    "    }\n",
-    "    \n",
-    "    # Exclude non-amino acid residues\n",
-    "    sequence = [\n",
-    "        residue for residue in chain \n",
-    "        if residue.get_resname().strip() in aa_dict\n",
-    "    ]\n",
-    "    \n",
-    "    random_scores = np.random.rand(len(sequence))\n",
-    "    result_str = \"\\n\".join(\n",
-    "        f\"{aa_dict[res.get_resname()]} {res.id[1]} {score:.2f}\" \n",
-    "        for res, score in zip(sequence, random_scores)\n",
-    "    )\n",
-    "    \n",
-    "    # Save the predictions to a file\n",
-    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
-    "    with open(prediction_file, \"w\") as f:\n",
-    "        f.write(result_str)\n",
-    "    \n",
-    "    return result_str, molecule(pdb_path, random_scores, segment), prediction_file\n",
-    "\n",
-    "def molecule(input_pdb, scores=None, segment='A'):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
-    "    \n",
-    "    # Prepare high-scoring residues script if scores are provided\n",
-    "    high_score_script = \"\"\n",
-    "    if scores is not None:\n",
-    "        high_score_script = \"\"\"\n",
-    "        // Reset all styles first\n",
-    "        viewer.getModel(0).setStyle({}, {});\n",
-    "        \n",
-    "        // Show only the selected chain\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\"}, \n",
-    "            { cartoon: {colorscheme:\"whiteCarbon\"} }\n",
-    "        );\n",
-    "        \n",
-    "        // Highlight high-scoring residues only for the selected chain\n",
-    "        let highScoreResidues = [%s];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": highScoreResidues}, \n",
-    "            {\"stick\": {\"color\": \"red\"}}\n",
-    "        );\n",
-    "\n",
-    "        // Highlight high-scoring residues only for the selected chain\n",
-    "        let highScoreResidues2 = [%s];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": highScoreResidues2}, \n",
-    "            {\"stick\": {\"color\": \"orange\"}}\n",
-    "        );\n",
-    "        \"\"\" % (segment, \n",
-    "               \", \".join(str(i+1) for i, score in enumerate(scores) if score > 0.8),\n",
-    "               segment,\n",
-    "              \", \".join(str(i+1) for i, score in enumerate(scores) if (score > 0.5) and (score < 0.8)),\n",
-    "               segment)\n",
-    "    \n",
-    "    html_content = f\"\"\"\n",
-    "    <!DOCTYPE html>\n",
-    "    <html>\n",
-    "    <head>    \n",
-    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
-    "        <style>\n",
-    "        .mol-container {{\n",
-    "            width: 100%;\n",
-    "            height: 700px;\n",
-    "            position: relative;\n",
-    "        }}\n",
-    "        </style>\n",
-    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
-    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
-    "    </head>\n",
-    "    <body>\n",
-    "        <div id=\"container\" class=\"mol-container\"></div>\n",
-    "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
-    "            $(document).ready(function () {{\n",
-    "                let element = $(\"#container\");\n",
-    "                let config = {{ backgroundColor: \"white\" }};\n",
-    "                let viewer = $3Dmol.createViewer(element, config);\n",
-    "                viewer.addModel(pdb, \"pdb\");\n",
-    "                \n",
-    "                // Reset all styles and show only selected chain\n",
-    "                viewer.getModel(0).setStyle(\n",
-    "                    {{\"chain\": \"{segment}\"}}, \n",
-    "                    {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }}\n",
-    "                );\n",
-    "                \n",
-    "                {high_score_script}\n",
-    "                \n",
-    "                // Add hover functionality\n",
-    "                viewer.setHoverable(\n",
-    "                    {{}}, \n",
-    "                    true, \n",
-    "                    function(atom, viewer, event, container) {{\n",
-    "                        if (!atom.label) {{\n",
-    "                            atom.label = viewer.addLabel(\n",
-    "                                atom.resn + \":\" +atom.resi + \":\" + atom.atom, \n",
-    "                                {{\n",
-    "                                    position: atom, \n",
-    "                                    backgroundColor: 'mintcream', \n",
-    "                                    fontColor: 'black',\n",
-    "                                    fontSize: 12,\n",
-    "                                    padding: 2\n",
-    "                                }}\n",
-    "                            );\n",
-    "                        }}\n",
-    "                    }},\n",
-    "                    function(atom, viewer) {{\n",
-    "                        if (atom.label) {{\n",
-    "                            viewer.removeLabel(atom.label);\n",
-    "                            delete atom.label;\n",
-    "                        }}\n",
-    "                    }}\n",
-    "                );\n",
-    "                \n",
-    "                viewer.zoomTo();\n",
-    "                viewer.render();\n",
-    "                viewer.zoom(0.8, 2000);\n",
-    "            }});\n",
-    "        </script>\n",
-    "    </body>\n",
-    "    </html>\n",
-    "    \"\"\"\n",
-    "    \n",
-    "    # Return the HTML content within an iframe safely encoded for special characters\n",
-    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
-    "\n",
-    "reps =    [\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"style\": \"cartoon\",\n",
-    "          \"color\": \"whiteCarbon\",\n",
-    "          \"residue_range\": \"\",\n",
-    "          \"around\": 0,\n",
-    "          \"byres\": False,\n",
-    "        }\n",
-    "    ]\n",
-    "\n",
-    "# Gradio UI\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
-    "\n",
-    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
-    "\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
-    "\n",
-    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
-    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
-    "    download_output = gr.File(label=\"Download Predictions\")\n",
-    "    \n",
-    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)\n",
-    "    \n",
-    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
-    "    \n",
-    "    gr.Markdown(\"## Examples\")\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            [\"2IWI\", \"A\"],\n",
-    "            [\"7RPZ\", \"B\"],\n",
-    "            [\"3TJN\", \"C\"]\n",
-    "        ],\n",
-    "        inputs=[pdb_input, segment_input],\n",
-    "        outputs=[predictions_output, molecule_output, download_output]\n",
-    "    )\n",
-    "\n",
-    "demo.launch()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "30f35243-852f-4771-9a4b-5cdd198552b5",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "* Running on local URL:  http://127.0.0.1:7865\n",
-      "\n",
-      "To create a public link, set `share=True` in `launch()`.\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<div><iframe src=\"http://127.0.0.1:7865/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/plain": []
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import gradio as gr\n",
-    "import requests\n",
-    "from Bio.PDB import PDBParser\n",
-    "import numpy as np\n",
-    "import os\n",
-    "from gradio_molecule3d import Molecule3D\n",
-    "\n",
-    "def read_mol(pdb_path):\n",
-    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
-    "    with open(pdb_path, 'r') as f:\n",
-    "        return f.read()\n",
-    "\n",
-    "def fetch_pdb(pdb_id):\n",
-    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
-    "    pdb_path = f'{pdb_id}.pdb'\n",
-    "    response = requests.get(pdb_url)\n",
-    "    if response.status_code == 200:\n",
-    "        with open(pdb_path, 'wb') as f:\n",
-    "            f.write(response.content)\n",
-    "        return pdb_path\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "def process_pdb(pdb_id, segment):\n",
-    "    pdb_path = fetch_pdb(pdb_id)\n",
-    "    if not pdb_path:\n",
-    "        return \"Failed to fetch PDB file\", None, None\n",
-    "    \n",
-    "    parser = PDBParser(QUIET=1)\n",
-    "    structure = parser.get_structure('protein', pdb_path)\n",
-    "    \n",
-    "    try:\n",
-    "        chain = structure[0][segment]\n",
-    "    except KeyError:\n",
-    "        return \"Invalid Chain ID\", None, None\n",
-    "    \n",
-    "    # Comprehensive amino acid mapping\n",
-    "    aa_dict = {\n",
-    "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
-    "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
-    "        'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',\n",
-    "        'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',\n",
-    "        'MSE': 'M', 'SEP': 'S', 'TPO': 'T', 'CSO': 'C', 'PTR': 'Y', 'HYP': 'P'\n",
-    "    }\n",
-    "    \n",
-    "    # Exclude non-amino acid residues and create a list of (resi, score) pairs\n",
-    "    sequence = [\n",
-    "        (res.id[1], res) for res in chain\n",
-    "        if res.get_resname().strip() in aa_dict\n",
-    "    ]\n",
-    "\n",
-    "    random_scores = np.random.rand(len(sequence))\n",
-    "    \n",
-    "    # Zip residues with scores to track the residue ID and score\n",
-    "    residue_scores = [(resi, score) for (resi, _), score in zip(sequence, random_scores)]\n",
-    "    \n",
-    "    result_str = \"\\n\".join(\n",
-    "        f\"{aa_dict[chain[resi].get_resname()]} {resi} {score:.2f}\"\n",
-    "        for resi, score in residue_scores\n",
-    "    )\n",
-    "    \n",
-    "    # Save the predictions to a file\n",
-    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
-    "    with open(prediction_file, \"w\") as f:\n",
-    "        f.write(result_str)\n",
-    "    \n",
-    "    return result_str, molecule(pdb_path, residue_scores, segment), prediction_file\n",
-    "\n",
-    "def molecule(input_pdb, residue_scores=None, segment='A'):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
-    "    \n",
-    "    # Prepare high-scoring residues script if scores are provided\n",
-    "    high_score_script = \"\"\n",
-    "    if residue_scores is not None:\n",
-    "        # Sort residues based on their scores\n",
-    "        high_score_residues = [resi for resi, score in residue_scores if score > 0.9]\n",
-    "        mid_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 0.9]\n",
-    "        \n",
-    "        high_score_script = \"\"\"\n",
-    "        // Reset all styles first\n",
-    "        viewer.getModel(0).setStyle({}, {});\n",
-    "        \n",
-    "        // Show only the selected chain\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\"}, \n",
-    "            { cartoon: {colorscheme:\"whiteCarbon\"} }\n",
-    "        );\n",
-    "        \n",
-    "        // Highlight high-scoring residues only for the selected chain\n",
-    "        let highScoreResidues = [%s];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": highScoreResidues}, \n",
-    "            {\"stick\": {\"color\": \"red\"}}\n",
-    "        );\n",
-    "\n",
-    "        // Highlight medium-scoring residues only for the selected chain\n",
-    "        let midScoreResidues = [%s];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": midScoreResidues}, \n",
-    "            {\"stick\": {\"color\": \"orange\"}}\n",
-    "        );\n",
-    "        \"\"\" % (segment, \n",
-    "               \", \".join(str(resi) for resi in high_score_residues),\n",
-    "               segment,\n",
-    "               \", \".join(str(resi) for resi in mid_score_residues),\n",
-    "               segment)\n",
-    "    \n",
-    "    html_content = f\"\"\"\n",
-    "    <!DOCTYPE html>\n",
-    "    <html>\n",
-    "    <head>    \n",
-    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
-    "        <style>\n",
-    "        .mol-container {{\n",
-    "            width: 100%;\n",
-    "            height: 700px;\n",
-    "            position: relative;\n",
-    "        }}\n",
-    "        </style>\n",
-    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
-    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
-    "    </head>\n",
-    "    <body>\n",
-    "        <div id=\"container\" class=\"mol-container\"></div>\n",
-    "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
-    "            $(document).ready(function () {{\n",
-    "                let element = $(\"#container\");\n",
-    "                let config = {{ backgroundColor: \"white\" }};\n",
-    "                let viewer = $3Dmol.createViewer(element, config);\n",
-    "                viewer.addModel(pdb, \"pdb\");\n",
-    "                \n",
-    "                // Reset all styles and show only selected chain\n",
-    "                viewer.getModel(0).setStyle(\n",
-    "                    {{\"chain\": \"{segment}\"}}, \n",
-    "                    {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }}\n",
-    "                );\n",
-    "                \n",
-    "                {high_score_script}\n",
-    "                \n",
-    "                // Add hover functionality\n",
-    "                viewer.setHoverable(\n",
-    "                    {{}}, \n",
-    "                    true, \n",
-    "                    function(atom, viewer, event, container) {{\n",
-    "                        if (!atom.label) {{\n",
-    "                            atom.label = viewer.addLabel(\n",
-    "                                atom.resn + \":\" +atom.resi + \":\" + atom.atom, \n",
-    "                                {{\n",
-    "                                    position: atom, \n",
-    "                                    backgroundColor: 'mintcream', \n",
-    "                                    fontColor: 'black',\n",
-    "                                    fontSize: 12,\n",
-    "                                    padding: 2\n",
-    "                                }}\n",
-    "                            );\n",
-    "                        }}\n",
-    "                    }},\n",
-    "                    function(atom, viewer) {{\n",
-    "                        if (atom.label) {{\n",
-    "                            viewer.removeLabel(atom.label);\n",
-    "                            delete atom.label;\n",
-    "                        }}\n",
-    "                    }}\n",
-    "                );\n",
-    "                \n",
-    "                viewer.zoomTo();\n",
-    "                viewer.render();\n",
-    "                viewer.zoom(0.8, 2000);\n",
-    "            }});\n",
-    "        </script>\n",
-    "    </body>\n",
-    "    </html>\n",
-    "    \"\"\"\n",
-    "    \n",
-    "    # Return the HTML content within an iframe safely encoded for special characters\n",
-    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
-    "\n",
-    "reps =    [\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"style\": \"cartoon\",\n",
-    "          \"color\": \"whiteCarbon\",\n",
-    "          \"residue_range\": \"\",\n",
-    "          \"around\": 0,\n",
-    "          \"byres\": False,\n",
-    "        }\n",
-    "    ]\n",
-    "\n",
-    "# Gradio UI\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
-    "\n",
-    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
-    "\n",
-    "    with gr.Row():\n",
-    "        #pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
-    "\n",
-    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
-    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
-    "    download_output = gr.File(label=\"Download Predictions\")\n",
-    "    \n",
-    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)\n",
-    "    \n",
-    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
-    "    \n",
-    "    gr.Markdown(\"## Examples\")\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            [\"7RPZ\", \"A\"],\n",
-    "            [\"2IWI\", \"B\"],\n",
-    "            [\"2F6V\", \"A\"]\n",
-    "        ],\n",
-    "        inputs=[pdb_input, segment_input],\n",
-    "        outputs=[predictions_output, molecule_output, download_output]\n",
-    "    )\n",
-    "\n",
-    "demo.launch()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6f17feec-0347-4f9d-acd4-ae681c3ed425",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "63201f38-adde-4b12-a8d3-f23474d045cf",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "5ccbf398-5ef2-4955-98db-99f904f8daa4",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4c61bac4-4f2e-4f4a-aa1f-30dca209747c",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import gradio as gr\n",
-    "import requests\n",
-    "from Bio.PDB import PDBParser\n",
-    "import numpy as np\n",
-    "import os\n",
-    "from gradio_molecule3d import Molecule3D\n",
-    "\n",
-    "\n",
-    "from model_loader import load_model\n",
-    "\n",
-    "import torch\n",
-    "import torch.nn as nn\n",
-    "import torch.nn.functional as F\n",
-    "from torch.utils.data import DataLoader\n",
-    "\n",
-    "import re\n",
-    "import pandas as pd\n",
-    "import copy\n",
-    "\n",
-    "import transformers, datasets\n",
-    "from transformers import AutoTokenizer\n",
-    "from transformers import DataCollatorForTokenClassification\n",
-    "\n",
-    "from datasets import Dataset\n",
-    "\n",
-    "from scipy.special import expit\n",
-    "\n",
-    "# Load model and move to device\n",
-    "checkpoint = 'ThorbenF/prot_t5_xl_uniref50'\n",
-    "max_length = 1500\n",
-    "model, tokenizer = load_model(checkpoint, max_length)\n",
-    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
-    "model.to(device)\n",
-    "model.eval()\n",
-    "\n",
-    "def normalize_scores(scores):\n",
-    "    min_score = np.min(scores)\n",
-    "    max_score = np.max(scores)\n",
-    "    return (scores - min_score) / (max_score - min_score) if max_score > min_score else scores\n",
-    "    \n",
-    "def read_mol(pdb_path):\n",
-    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
-    "    with open(pdb_path, 'r') as f:\n",
-    "        return f.read()\n",
-    "\n",
-    "def fetch_pdb(pdb_id):\n",
-    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
-    "    pdb_path = f'{pdb_id}.pdb'\n",
-    "    response = requests.get(pdb_url)\n",
-    "    if response.status_code == 200:\n",
-    "        with open(pdb_path, 'wb') as f:\n",
-    "            f.write(response.content)\n",
-    "        return pdb_path\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "def process_pdb(pdb_id, segment):\n",
-    "    pdb_path = fetch_pdb(pdb_id)\n",
-    "    if not pdb_path:\n",
-    "        return \"Failed to fetch PDB file\", None, None\n",
-    "    \n",
-    "    parser = PDBParser(QUIET=1)\n",
-    "    structure = parser.get_structure('protein', pdb_path)\n",
-    "    \n",
-    "    try:\n",
-    "        chain = structure[0][segment]\n",
-    "    except KeyError:\n",
-    "        return \"Invalid Chain ID\", None, None\n",
-    "    \n",
-    "    \n",
-    "    aa_dict = {\n",
-    "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
-    "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
-    "        'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',\n",
-    "        'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',\n",
-    "        'MSE': 'M', 'SEP': 'S', 'TPO': 'T', 'CSO': 'C', 'PTR': 'Y', 'HYP': 'P'\n",
-    "    }\n",
-    "    \n",
-    "    # Exclude non-amino acid residues\n",
-    "    sequence = \"\".join(\n",
-    "        aa_dict[residue.get_resname().strip()] \n",
-    "        for residue in chain \n",
-    "        if residue.get_resname().strip() in aa_dict\n",
-    "    )\n",
-    "    sequence2 = [\n",
-    "        (res.id[1], res) for res in chain\n",
-    "        if res.get_resname().strip() in aa_dict\n",
-    "    ]\n",
-    "    \n",
-    "    # Prepare input for model prediction\n",
-    "    input_ids = tokenizer(\" \".join(sequence), return_tensors=\"pt\").input_ids.to(device)\n",
-    "    with torch.no_grad():\n",
-    "        outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()\n",
-    "\n",
-    "    # Calculate scores and normalize them\n",
-    "    scores = expit(outputs[:, 1] - outputs[:, 0])\n",
-    "    normalized_scores = normalize_scores(scores)\n",
-    "\n",
-    "    # Zip residues with scores to track the residue ID and score\n",
-    "    residue_scores = [(resi, score) for (resi, _), score in zip(sequence2, normalized_scores)]\n",
-    "    \n",
-    "    result_str = \"\\n\".join([\n",
-    "        f\"{res.get_resname()} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}\" \n",
-    "        for i, res in enumerate(chain) if res.get_resname().strip() in aa_dict\n",
-    "    ])\n",
-    "    \n",
-    "    # Save the predictions to a file\n",
-    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
-    "    with open(prediction_file, \"w\") as f:\n",
-    "        f.write(result_str)\n",
-    "    \n",
-    "    return result_str, molecule(pdb_path, residue_scores, segment), prediction_file\n",
-    "\n",
-    "def molecule(input_pdb, residue_scores=None, segment='A'):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
-    "    \n",
-    "    # Prepare high-scoring residues script if scores are provided\n",
-    "    high_score_script = \"\"\n",
-    "    if residue_scores is not None:\n",
-    "        # Sort residues based on their scores\n",
-    "        high_score_residues = [resi for resi, score in residue_scores if score > 0.75]\n",
-    "        mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75]\n",
-    "        \n",
-    "        high_score_script = \"\"\"\n",
-    "        // Reset all styles first\n",
-    "        viewer.getModel(0).setStyle({}, {});\n",
-    "        \n",
-    "        // Show only the selected chain\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\"}, \n",
-    "            { cartoon: {colorscheme:\"whiteCarbon\"} }\n",
-    "        );\n",
-    "        \n",
-    "        // Highlight high-scoring residues only for the selected chain\n",
-    "        let highScoreResidues = [%s];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": highScoreResidues}, \n",
-    "            {\"stick\": {\"color\": \"red\"}}\n",
-    "        );\n",
-    "\n",
-    "        // Highlight medium-scoring residues only for the selected chain\n",
-    "        let midScoreResidues = [%s];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": midScoreResidues}, \n",
-    "            {\"stick\": {\"color\": \"orange\"}}\n",
-    "        );\n",
-    "        \"\"\" % (segment, \n",
-    "               \", \".join(str(resi) for resi in high_score_residues),\n",
-    "               segment,\n",
-    "               \", \".join(str(resi) for resi in mid_score_residues),\n",
-    "               segment)\n",
-    "    \n",
-    "    html_content = f\"\"\"\n",
-    "    <!DOCTYPE html>\n",
-    "    <html>\n",
-    "    <head>    \n",
-    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
-    "        <style>\n",
-    "        .mol-container {{\n",
-    "            width: 100%;\n",
-    "            height: 700px;\n",
-    "            position: relative;\n",
-    "        }}\n",
-    "        </style>\n",
-    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
-    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
-    "    </head>\n",
-    "    <body>\n",
-    "        <div id=\"container\" class=\"mol-container\"></div>\n",
-    "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
-    "            $(document).ready(function () {{\n",
-    "                let element = $(\"#container\");\n",
-    "                let config = {{ backgroundColor: \"white\" }};\n",
-    "                let viewer = $3Dmol.createViewer(element, config);\n",
-    "                viewer.addModel(pdb, \"pdb\");\n",
-    "                \n",
-    "                // Reset all styles and show only selected chain\n",
-    "                viewer.getModel(0).setStyle(\n",
-    "                    {{\"chain\": \"{segment}\"}}, \n",
-    "                    {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }}\n",
-    "                );\n",
-    "                \n",
-    "                {high_score_script}\n",
-    "                \n",
-    "                // Add hover functionality\n",
-    "                viewer.setHoverable(\n",
-    "                    {{}}, \n",
-    "                    true, \n",
-    "                    function(atom, viewer, event, container) {{\n",
-    "                        if (!atom.label) {{\n",
-    "                            atom.label = viewer.addLabel(\n",
-    "                                atom.resn + \":\" +atom.resi + \":\" + atom.atom, \n",
-    "                                {{\n",
-    "                                    position: atom, \n",
-    "                                    backgroundColor: 'mintcream', \n",
-    "                                    fontColor: 'black',\n",
-    "                                    fontSize: 12,\n",
-    "                                    padding: 2\n",
-    "                                }}\n",
-    "                            );\n",
-    "                        }}\n",
-    "                    }},\n",
-    "                    function(atom, viewer) {{\n",
-    "                        if (atom.label) {{\n",
-    "                            viewer.removeLabel(atom.label);\n",
-    "                            delete atom.label;\n",
-    "                        }}\n",
-    "                    }}\n",
-    "                );\n",
-    "                \n",
-    "                viewer.zoomTo();\n",
-    "                viewer.render();\n",
-    "                viewer.zoom(0.8, 2000);\n",
-    "            }});\n",
-    "        </script>\n",
-    "    </body>\n",
-    "    </html>\n",
-    "    \"\"\"\n",
-    "    \n",
-    "    # Return the HTML content within an iframe safely encoded for special characters\n",
-    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
-    "\n",
-    "reps =    [\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"style\": \"cartoon\",\n",
-    "          \"color\": \"whiteCarbon\",\n",
-    "          \"residue_range\": \"\",\n",
-    "          \"around\": 0,\n",
-    "          \"byres\": False,\n",
-    "        }\n",
-    "    ]\n",
-    "\n",
-    "# Gradio UI\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
-    "\n",
-    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
-    "\n",
-    "    with gr.Row():\n",
-    "        #pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
-    "\n",
-    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
-    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
-    "    download_output = gr.File(label=\"Download Predictions\")\n",
-    "    \n",
-    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)\n",
-    "    \n",
-    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
-    "    \n",
-    "    gr.Markdown(\"## Examples\")\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            [\"7RPZ\", \"A\"],\n",
-    "            [\"2IWI\", \"B\"],\n",
-    "            [\"2F6V\", \"A\"]\n",
-    "        ],\n",
-    "        inputs=[pdb_input, segment_input],\n",
-    "        outputs=[predictions_output, molecule_output, download_output]\n",
-    "    )\n",
-    "\n",
-    "demo.launch(share=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "b61d06ec-a4ee-4f65-925f-d2688730416a",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4d67d69f-1f53-4bcc-8905-8d29384c4e20",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import gradio as gr\n",
-    "import requests\n",
-    "from Bio.PDB import PDBParser\n",
-    "import numpy as np\n",
-    "import os\n",
-    "from gradio_molecule3d import Molecule3D\n",
-    "\n",
-    "\n",
-    "from model_loader import load_model\n",
-    "\n",
-    "import torch\n",
-    "import torch.nn as nn\n",
-    "import torch.nn.functional as F\n",
-    "from torch.utils.data import DataLoader\n",
-    "\n",
-    "import re\n",
-    "import pandas as pd\n",
-    "import copy\n",
-    "\n",
-    "import transformers, datasets\n",
-    "from transformers import AutoTokenizer\n",
-    "from transformers import DataCollatorForTokenClassification\n",
-    "\n",
-    "from datasets import Dataset\n",
-    "\n",
-    "from scipy.special import expit\n",
-    "\n",
-    "# Load model and move to device\n",
-    "checkpoint = 'ThorbenF/prot_t5_xl_uniref50'\n",
-    "max_length = 1500\n",
-    "model, tokenizer = load_model(checkpoint, max_length)\n",
-    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
-    "model.to(device)\n",
-    "model.eval()\n",
-    "\n",
-    "def normalize_scores(scores):\n",
-    "    min_score = np.min(scores)\n",
-    "    max_score = np.max(scores)\n",
-    "    return (scores - min_score) / (max_score - min_score) if max_score > min_score else scores\n",
-    "    \n",
-    "def read_mol(pdb_path):\n",
-    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
-    "    with open(pdb_path, 'r') as f:\n",
-    "        return f.read()\n",
-    "\n",
-    "def fetch_pdb(pdb_id):\n",
-    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
-    "    pdb_path = f'{pdb_id}.pdb'\n",
-    "    response = requests.get(pdb_url)\n",
-    "    if response.status_code == 200:\n",
-    "        with open(pdb_path, 'wb') as f:\n",
-    "            f.write(response.content)\n",
-    "        return pdb_path\n",
-    "    else:\n",
-    "        return None\n",
-    "\n",
-    "def process_pdb(pdb_id, segment):\n",
-    "    pdb_path = fetch_pdb(pdb_id)\n",
-    "    if not pdb_path:\n",
-    "        return \"Failed to fetch PDB file\", None, None\n",
-    "    \n",
-    "    parser = PDBParser(QUIET=1)\n",
-    "    structure = parser.get_structure('protein', pdb_path)\n",
-    "    \n",
-    "    try:\n",
-    "        chain = structure[0][segment]\n",
-    "    except KeyError:\n",
-    "        return \"Invalid Chain ID\", None, None\n",
-    "    \n",
-    "    \n",
-    "    aa_dict = {\n",
-    "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
-    "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
-    "        'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',\n",
-    "        'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',\n",
-    "        'MSE': 'M', 'SEP': 'S', 'TPO': 'T', 'CSO': 'C', 'PTR': 'Y', 'HYP': 'P'\n",
-    "    }\n",
-    "    \n",
-    "    # Exclude non-amino acid residues\n",
-    "    sequence = \"\".join(\n",
-    "        aa_dict[residue.get_resname().strip()] \n",
-    "        for residue in chain \n",
-    "        if residue.get_resname().strip() in aa_dict\n",
-    "    )\n",
-    "    sequence2 = [\n",
-    "        (res.id[1], res) for res in chain\n",
-    "        if res.get_resname().strip() in aa_dict\n",
-    "    ]\n",
-    "    \n",
-    "    # Prepare input for model prediction\n",
-    "    input_ids = tokenizer(\" \".join(sequence), return_tensors=\"pt\").input_ids.to(device)\n",
-    "    with torch.no_grad():\n",
-    "        outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()\n",
-    "\n",
-    "    # Calculate scores and normalize them\n",
-    "    scores = expit(outputs[:, 1] - outputs[:, 0])\n",
-    "    normalized_scores = normalize_scores(scores)\n",
-    "\n",
-    "    # Zip residues with scores to track the residue ID and score\n",
-    "    residue_scores = [(resi, score) for (resi, _), score in zip(sequence2, normalized_scores)]\n",
-    "    \n",
-    "    result_str = \"\\n\".join([\n",
-    "        f\"{res.get_resname()} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}\" \n",
-    "        for i, res in enumerate(chain) if res.get_resname().strip() in aa_dict\n",
-    "    ])\n",
-    "    \n",
-    "    # Save the predictions to a file\n",
-    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
-    "    with open(prediction_file, \"w\") as f:\n",
-    "        f.write(result_str)\n",
-    "    \n",
-    "    return result_str, molecule(pdb_path, residue_scores, segment), prediction_file\n",
-    "\n",
-    "def molecule(input_pdb, residue_scores=None, segment='A'):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
-    "    \n",
-    "    # Prepare high-scoring residues script if scores are provided\n",
-    "    high_score_script = \"\"\n",
-    "    if residue_scores is not None:\n",
-    "        # Sort residues based on their scores\n",
-    "        high_score_residues = [resi for resi, score in residue_scores if score > 0.75]\n",
-    "        mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75]\n",
-    "        \n",
-    "        high_score_script = \"\"\"\n",
-    "        // Reset all styles first\n",
-    "        viewer.getModel(0).setStyle({}, {});\n",
-    "        \n",
-    "        // Show only the selected chain\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\"}, \n",
-    "            { cartoon: {colorscheme:\"whiteCarbon\"} }\n",
-    "        );\n",
-    "        \n",
-    "        // Highlight high-scoring residues only for the selected chain\n",
-    "        let highScoreResidues = [%s];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": highScoreResidues}, \n",
-    "            {\"stick\": {\"color\": \"red\"}}\n",
-    "        );\n",
-    "\n",
-    "        // Highlight medium-scoring residues only for the selected chain\n",
-    "        let midScoreResidues = [%s];\n",
-    "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": midScoreResidues}, \n",
-    "            {\"stick\": {\"color\": \"orange\"}}\n",
-    "        );\n",
-    "        \"\"\" % (segment, \n",
-    "               \", \".join(str(resi) for resi in high_score_residues),\n",
-    "               segment,\n",
-    "               \", \".join(str(resi) for resi in mid_score_residues),\n",
-    "               segment)\n",
-    "    \n",
-    "    html_content = f\"\"\"\n",
-    "    <!DOCTYPE html>\n",
-    "    <html>\n",
-    "    <head>    \n",
-    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
-    "        <style>\n",
-    "        .mol-container {{\n",
-    "            width: 100%;\n",
-    "            height: 700px;\n",
-    "            position: relative;\n",
-    "        }}\n",
-    "        </style>\n",
-    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
-    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
-    "    </head>\n",
-    "    <body>\n",
-    "        <div id=\"container\" class=\"mol-container\"></div>\n",
-    "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
-    "            $(document).ready(function () {{\n",
-    "                let element = $(\"#container\");\n",
-    "                let config = {{ backgroundColor: \"white\" }};\n",
-    "                let viewer = $3Dmol.createViewer(element, config);\n",
-    "                viewer.addModel(pdb, \"pdb\");\n",
-    "                \n",
-    "                // Reset all styles and show only selected chain\n",
-    "                viewer.getModel(0).setStyle(\n",
-    "                    {{\"chain\": \"{segment}\"}}, \n",
-    "                    {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }}\n",
-    "                );\n",
-    "                \n",
-    "                {high_score_script}\n",
-    "                \n",
-    "                // Add hover functionality\n",
-    "                viewer.setHoverable(\n",
-    "                    {{}}, \n",
-    "                    true, \n",
-    "                    function(atom, viewer, event, container) {{\n",
-    "                        if (!atom.label) {{\n",
-    "                            atom.label = viewer.addLabel(\n",
-    "                                atom.resn + \":\" +atom.resi + \":\" + atom.atom, \n",
-    "                                {{\n",
-    "                                    position: atom, \n",
-    "                                    backgroundColor: 'mintcream', \n",
-    "                                    fontColor: 'black',\n",
-    "                                    fontSize: 12,\n",
-    "                                    padding: 2\n",
-    "                                }}\n",
-    "                            );\n",
-    "                        }}\n",
-    "                    }},\n",
-    "                    function(atom, viewer) {{\n",
-    "                        if (atom.label) {{\n",
-    "                            viewer.removeLabel(atom.label);\n",
-    "                            delete atom.label;\n",
-    "                        }}\n",
-    "                    }}\n",
-    "                );\n",
-    "                \n",
-    "                viewer.zoomTo();\n",
-    "                viewer.render();\n",
-    "                viewer.zoom(0.8, 2000);\n",
-    "            }});\n",
-    "        </script>\n",
-    "    </body>\n",
-    "    </html>\n",
-    "    \"\"\"\n",
-    "    \n",
-    "    # Return the HTML content within an iframe safely encoded for special characters\n",
-    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
-    "\n",
-    "reps =    [\n",
-    "        {\n",
-    "          \"model\": 0,\n",
-    "          \"style\": \"cartoon\",\n",
-    "          \"color\": \"whiteCarbon\",\n",
-    "          \"residue_range\": \"\",\n",
-    "          \"around\": 0,\n",
-    "          \"byres\": False,\n",
-    "        }\n",
-    "    ]\n",
-    "\n",
-    "# Gradio UI\n",
-    "with gr.Blocks() as demo:\n",
-    "    gr.Markdown(\"# Protein Binding Site Prediction\")\n",
-    "    with gr.Row():\n",
-    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
-    "\n",
-    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
-    "\n",
-    "    with gr.Row():\n",
-    "        #pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
-    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
-    "        prediction_btn = gr.Button(\"Predict Binding Site\")\n",
-    "\n",
-    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
-    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
-    "    download_output = gr.File(label=\"Download Predictions\")\n",
-    "    \n",
-    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)\n",
-    "    \n",
-    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
-    "    \n",
-    "    gr.Markdown(\"## Examples\")\n",
-    "    gr.Examples(\n",
-    "        examples=[\n",
-    "            [\"7RPZ\", \"A\"],\n",
-    "            [\"2IWI\", \"B\"],\n",
-    "            [\"2F6V\", \"A\"]\n",
-    "        ],\n",
-    "        inputs=[pdb_input, segment_input],\n",
-    "        outputs=[predictions_output, molecule_output, download_output]\n",
-    "    )\n",
-    "\n",
-    "demo.launch(share=True)"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python (LLM)",
-   "language": "python",
-   "name": "llm"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.12.7"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}