Update

.ipynb_checkpoints/app-checkpoint.py

@@ -1,4 +1,11 @@
 import gradio as gr
+import requests
+from Bio.PDB import PDBParser
+import numpy as np
+import os
+from gradio_molecule3d import Molecule3D
+
+
 from model_loader import load_model
 
 import torch
@@ -7,8 +14,6 @@ import torch.nn.functional as F
 from torch.utils.data import DataLoader
 
 import re
-import numpy as np
-import os
 import pandas as pd
 import copy
 
@@ -20,18 +25,6 @@ from datasets import Dataset
 
 from scipy.special import expit
 
-import requests
-
-from gradio_molecule3d import Molecule3D
-
-# Biopython imports
-from Bio.PDB import PDBParser, Select, PDBIO
-from Bio.PDB.DSSP import DSSP
-from Bio.PDB import PDBList
-
-from matplotlib import cm  # For color mapping
-from matplotlib.colors import Normalize
-
 # Load model and move to device
 checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
 max_length = 1500
@@ -40,23 +33,26 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)
 model.eval()
 
-# Function to fetch a PDB file
+def normalize_scores(scores):
+    min_score = np.min(scores)
+    max_score = np.max(scores)
+    return (scores - min_score) / (max_score - min_score) if max_score > min_score else scores
+    
+def read_mol(pdb_path):
+    """Read PDB file and return its content as a string"""
+    with open(pdb_path, 'r') as f:
+        return f.read()
+
 def fetch_pdb(pdb_id):
     pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'
-    pdb_path = f'pdb_files/{pdb_id}.pdb'
-    os.makedirs('pdb_files', exist_ok=True)
+    pdb_path = f'{pdb_id}.pdb'
     response = requests.get(pdb_url)
     if response.status_code == 200:
         with open(pdb_path, 'wb') as f:
             f.write(response.content)
         return pdb_path
-    return None
-
-
-def normalize_scores(scores):
-    min_score = np.min(scores)
-    max_score = np.max(scores)
-    return (scores - min_score) / (max_score - min_score) if max_score > min_score else scores
+    else:
+        return None
 
 def process_pdb(pdb_id, segment):
     pdb_path = fetch_pdb(pdb_id)
@@ -65,7 +61,11 @@ def process_pdb(pdb_id, segment):
     
     parser = PDBParser(QUIET=1)
     structure = parser.get_structure('protein', pdb_path)
-    chain = structure[0][segment]
+    
+    try:
+        chain = structure[0][segment]
+    except KeyError:
+        return "Invalid Chain ID", None, None
     
     # Comprehensive amino acid mapping
     aa_dict = {
@@ -77,11 +77,10 @@ def process_pdb(pdb_id, segment):
     }
     
     # Exclude non-amino acid residues
-    sequence = "".join(
-        aa_dict[residue.get_resname().strip()] 
-        for residue in chain 
+    sequence = [
+        residue for residue in chain 
         if residue.get_resname().strip() in aa_dict
-    )
+    ]
     
     # Prepare input for model prediction
     input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
@@ -91,54 +90,166 @@ def process_pdb(pdb_id, segment):
     # Calculate scores and normalize them
     scores = expit(outputs[:, 1] - outputs[:, 0])
     normalized_scores = normalize_scores(scores)
+
+    result_str = "\n".join(
+        f"{aa_dict[res.get_resname()]} {res.id[1]} {score:.2f}" 
+        for res, score in zip(sequence, normalized_scores)
+    )
     
-    # Prepare the result string, including only amino acid residues
-    result_str = "\n".join([
-        f"{res.get_resname()} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}" 
-        for i, res in enumerate(chain) if res.get_resname().strip() in aa_dict
-    ])
-    
-    # Save predictions to file
-    with open(f"{pdb_id}_predictions.txt", "w") as f:
+    # Save the predictions to a file
+    prediction_file = f"{pdb_id}_predictions.txt"
+    with open(prediction_file, "w") as f:
         f.write(result_str)
     
-    return result_str, pdb_path, f"{pdb_id}_predictions.txt"
+    return result_str, molecule(pdb_path, random_scores, segment), prediction_file
+
+def molecule(input_pdb, scores=None, segment='A'):
+    mol = read_mol(input_pdb)  # Read PDB file content
+    
+    # Prepare high-scoring residues script if scores are provided
+    high_score_script = ""
+    if scores is not None:
+        high_score_script = """
+        // Reset all styles first
+        viewer.getModel(0).setStyle({}, {});
+        
+        // Show only the selected chain
+        viewer.getModel(0).setStyle(
+            {"chain": "%s"}, 
+            { cartoon: {colorscheme:"whiteCarbon"} }
+        );
+        
+        // Highlight high-scoring residues only for the selected chain
+        let highScoreResidues = [%s];
+        viewer.getModel(0).setStyle(
+            {"chain": "%s", "resi": highScoreResidues}, 
+            {"stick": {"color": "red"}}
+        );
 
-reps = [{"model": 0, "style": "cartoon", "color": "spectrum"}]
+        // Highlight high-scoring residues only for the selected chain
+        let highScoreResidues2 = [%s];
+        viewer.getModel(0).setStyle(
+            {"chain": "%s", "resi": highScoreResidues2}, 
+            {"stick": {"color": "orange"}}
+        );
+        """ % (segment, 
+               ", ".join(str(i+1) for i, score in enumerate(scores) if score > 0.8),
+               segment,
+              ", ".join(str(i+1) for i, score in enumerate(scores) if (score > 0.5) and (score < 0.8)),
+               segment)
+    
+    html_content = f"""
+    <!DOCTYPE html>
+    <html>
+    <head>    
+        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+        <style>
+        .mol-container {{
+            width: 100%;
+            height: 700px;
+            position: relative;
+        }}
+        </style>
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js"></script>
+        <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+    </head>
+    <body>
+        <div id="container" class="mol-container"></div>
+        <script>
+            let pdb = `{mol}`; // Use template literal to properly escape PDB content
+            $(document).ready(function () {{
+                let element = $("#container");
+                let config = {{ backgroundColor: "white" }};
+                let viewer = $3Dmol.createViewer(element, config);
+                viewer.addModel(pdb, "pdb");
+                
+                // Reset all styles and show only selected chain
+                viewer.getModel(0).setStyle(
+                    {{"chain": "{segment}"}}, 
+                    {{ cartoon: {{ colorscheme:"whiteCarbon" }} }}
+                );
+                
+                {high_score_script}
+                
+                // Add hover functionality
+                viewer.setHoverable(
+                    {{}}, 
+                    true, 
+                    function(atom, viewer, event, container) {{
+                        if (!atom.label) {{
+                            atom.label = viewer.addLabel(
+                                atom.resn + ":" + atom.atom, 
+                                {{
+                                    position: atom, 
+                                    backgroundColor: 'mintcream', 
+                                    fontColor: 'black',
+                                    fontSize: 12,
+                                    padding: 2
+                                }}
+                            );
+                        }}
+                    }},
+                    function(atom, viewer) {{
+                        if (atom.label) {{
+                            viewer.removeLabel(atom.label);
+                            delete atom.label;
+                        }}
+                    }}
+                );
+                
+                viewer.zoomTo();
+                viewer.render();
+                viewer.zoom(0.8, 2000);
+            }});
+        </script>
+    </body>
+    </html>
+    """
+    
+    # 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>'
+
+reps =    [
+        {
+          "model": 0,
+          "style": "cartoon",
+          "color": "whiteCarbon",
+          "residue_range": "",
+          "around": 0,
+          "byres": False,
+        }
+    ]
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# Protein Binding Site Prediction")
+    gr.Markdown("# Protein Binding Site Prediction (Random Scores)")
+    with gr.Row():
+        pdb_input = gr.Textbox(value="2IWI", label="PDB ID", placeholder="Enter PDB ID here...")
+        visualize_btn = gr.Button("Visualize Structure")
+
+    molecule_output2 = Molecule3D(label="Protein Structure", reps=reps)
 
     with gr.Row():
-        pdb_input = gr.Textbox(value="2IWI",
-                label="PDB ID",
-                placeholder="Enter PDB ID here...")
-        segment_input = gr.Textbox(value="A",
-                label="Chain ID (Segment)",
-                placeholder="Enter Chain ID here...")
-        visualize_btn = gr.Button("Visualize Sructure")
-        prediction_btn = gr.Button("Predict Ligand Binding Site")
-
-    molecule_output = Molecule3D(label="Protein Structure", reps=reps)
+        pdb_input = gr.Textbox(value="2IWI", label="PDB ID", placeholder="Enter PDB ID here...")
+        segment_input = gr.Textbox(value="A", label="Chain ID", placeholder="Enter Chain ID here...")
+        prediction_btn = gr.Button("Predict Random Binding Site Scores")
+
+    molecule_output = gr.HTML(label="Protein Structure")
     predictions_output = gr.Textbox(label="Binding Site Predictions")
     download_output = gr.File(label="Download Predictions")
-
-    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output)
-    prediction_btn.click(
-        process_pdb, 
-        inputs=[pdb_input, segment_input], 
-        outputs=[predictions_output, molecule_output, download_output]
-    )
-
+    
+    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)
+    
+    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])
+    
     gr.Markdown("## Examples")
     gr.Examples(
         examples=[
-            ["2IWI"],
-            ["7RPZ"],
-            ["3TJN"]
+            ["2IWI", "A"],
+            ["7RPZ", "B"],
+            ["3TJN", "C"]
         ],
-        inputs=[pdb_input, segment_input], 
+        inputs=[pdb_input, segment_input],
         outputs=[predictions_output, molecule_output, download_output]
     )
 

.ipynb_checkpoints/requirements-checkpoint.txt

@@ -10,5 +10,4 @@ sentencepiece
 huggingface_hub>=0.15.0
 requests
 gradio_molecule3d
-biopython>=1.81
-matplotlib
+biopython>=1.81

.ipynb_checkpoints/test-checkpoint.ipynb

@@ -0,0 +1,452 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "d2208d17-47b6-4ff1-b6b6-ba09a9d490c7",
+   "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": 5,
+     "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": "code",
+   "execution_count": 4,
+   "id": "a1088e14-f09c-48ff-8632-cc4685306d7c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on local URL:  http://127.0.0.1:7863\n",
+      "\n",
+      "To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7863/\" 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"
+    }
+   ],
+   "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": "code",
+   "execution_count": 8,
+   "id": "cdf7fd26-0464-40d9-9107-71c29dbcaef8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on local URL:  http://127.0.0.1:7867\n",
+      "\n",
+      "To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7867/\" 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": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/var/folders/tm/ym2tckv54b96ws82y3b7cqhh0000gn/T/ipykernel_11794/4072855226.py:39: MatplotlibDeprecationWarning: The get_cmap function was deprecated in Matplotlib 3.7 and will be removed in 3.11. Use ``matplotlib.colormaps[name]`` or ``matplotlib.colormaps.get_cmap()`` or ``pyplot.get_cmap()`` instead.\n",
+      "  colors = [cm.get_cmap('coolwarm')(score)[:3] for score in normalized_scores]\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/Users/thorben_froehlking/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/queueing.py\", line 622, in process_events\n",
+      "    response = await route_utils.call_process_api(\n",
+      "               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/Users/thorben_froehlking/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/route_utils.py\", line 323, in call_process_api\n",
+      "    output = await app.get_blocks().process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/Users/thorben_froehlking/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 2024, in process_api\n",
+      "    data = await self.postprocess_data(block_fn, result[\"prediction\"], state)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/Users/thorben_froehlking/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 1830, in postprocess_data\n",
+      "    prediction_value = block.postprocess(prediction_value)\n",
+      "                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/Users/thorben_froehlking/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio_molecule3d/molecule3d.py\", line 210, in postprocess\n",
+      "    orig_name=Path(file).name,\n",
+      "              ^^^^^^^^^^\n",
+      "  File \"/Users/thorben_froehlking/anaconda3/envs/LLM/lib/python3.12/pathlib.py\", line 1162, in __init__\n",
+      "    super().__init__(*args)\n",
+      "  File \"/Users/thorben_froehlking/anaconda3/envs/LLM/lib/python3.12/pathlib.py\", line 373, in __init__\n",
+      "    raise TypeError(\n",
+      "TypeError: argument should be a str or an os.PathLike object where __fspath__ returns a str, not 'dict'\n"
+     ]
+    }
+   ],
+   "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",
+    "from matplotlib import cm\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, 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, 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",
+    "    # Normalize scores for coloring (0 = blue, 1 = red)\n",
+    "    normalized_scores = (random_scores - np.min(random_scores)) / (np.max(random_scores) - np.min(random_scores))\n",
+    "    colors = [cm.get_cmap('coolwarm')(score)[:3] for score in normalized_scores]\n",
+    "    hex_colors = [f'#{int(r*255):02x}{int(g*255):02x}{int(b*255):02x}' for r, g, b in colors]\n",
+    "\n",
+    "    # Result string and representation\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",
+    "    # Representation for the protein structure\n",
+    "    reps = [\n",
+    "        {\n",
+    "            \"model\": 0,\n",
+    "            \"style\": \"cartoon\",\n",
+    "            \"color\": \"whiteCarbon\"\n",
+    "        }\n",
+    "    ] + [\n",
+    "        {\n",
+    "            \"model\": 0,\n",
+    "            \"style\": \"cartoon\",\n",
+    "            \"residue_index\": i,\n",
+    "            \"color\": color\n",
+    "        }\n",
+    "        for i, color in enumerate(hex_colors)\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, reps, prediction_file\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",
+    "    prediction_btn.click(\n",
+    "        fn=process_pdb,\n",
+    "        inputs=[pdb_input, segment_input],\n",
+    "        outputs=[predictions_output, molecule_output, download_output]\n",
+    "    )\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": "ee215c16-a1fb-450f-bb93-37aaee6fb3f1",
+   "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
+}

.ipynb_checkpoints/test2-checkpoint.ipynb

@@ -0,0 +1,1193 @@
+{
+ "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": 11,
+   "id": "d62be1b5-762e-4b69-aed4-e4ba2a44482f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on local URL:  http://127.0.0.1:7867\n",
+      "\n",
+      "To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7867/\" 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": 11,
+     "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.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": "30f35243-852f-4771-9a4b-5cdd198552b5",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5eca6754-4aa1-463f-881a-25d2a0d6bb5b",
+   "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",
+    "    # 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",
+    "    # 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",
+    "    result_str = \"\\n\".join(\n",
+    "        f\"{aa_dict[res.get_resname()]} {res.id[1]} {score:.2f}\" \n",
+    "        for res, score in zip(sequence, normalized_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.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": "95046d1c-ec7c-4e3e-8a98-1802cb09a25b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a37cbe6f-d57f-41e5-8ae1-38258da39d47",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import gradio as gr\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 numpy as np\n",
+    "import os\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",
+    "import requests\n",
+    "\n",
+    "from gradio_molecule3d import Molecule3D\n",
+    "\n",
+    "# Biopython imports\n",
+    "from Bio.PDB import PDBParser, Select, PDBIO\n",
+    "from Bio.PDB.DSSP import DSSP\n",
+    "from Bio.PDB import PDBList\n",
+    "\n",
+    "from matplotlib import cm  # For color mapping\n",
+    "from matplotlib.colors import Normalize\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",
+    "# Function to fetch a PDB file\n",
+    "def fetch_pdb(pdb_id):\n",
+    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
+    "    pdb_path = f'pdb_files/{pdb_id}.pdb'\n",
+    "    os.makedirs('pdb_files', exist_ok=True)\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",
+    "    return None\n",
+    "\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 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",
+    "    chain = structure[0][segment]\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 = \"\".join(\n",
+    "        aa_dict[residue.get_resname().strip()] \n",
+    "        for residue in chain \n",
+    "        if residue.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",
+    "    # Prepare the result string, including only amino acid residues\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 predictions to file\n",
+    "    with open(f\"{pdb_id}_predictions.txt\", \"w\") as f:\n",
+    "        f.write(result_str)\n",
+    "    \n",
+    "    return result_str, pdb_path, f\"{pdb_id}_predictions.txt\"\n",
+    "\n",
+    "reps = [{\"model\": 0, \"style\": \"cartoon\", \"color\": \"spectrum\"}]\n",
+    "\n",
+    "# Gradio UI\n",
+    "with gr.Blocks() as demo:\n",
+    "    gr.Markdown(\"# Protein Binding Site Prediction\")\n",
+    "\n",
+    "    with gr.Row():\n",
+    "        pdb_input = gr.Textbox(value=\"2IWI\",\n",
+    "                label=\"PDB ID\",\n",
+    "                placeholder=\"Enter PDB ID here...\")\n",
+    "        segment_input = gr.Textbox(value=\"A\",\n",
+    "                label=\"Chain ID (Segment)\",\n",
+    "                placeholder=\"Enter Chain ID here...\")\n",
+    "        visualize_btn = gr.Button(\"Visualize Sructure\")\n",
+    "        prediction_btn = gr.Button(\"Predict Ligand Binding Site\")\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(\n",
+    "        process_pdb, \n",
+    "        inputs=[pdb_input, segment_input], \n",
+    "        outputs=[predictions_output, molecule_output, download_output]\n",
+    "    )\n",
+    "\n",
+    "    gr.Markdown(\"## Examples\")\n",
+    "    gr.Examples(\n",
+    "        examples=[\n",
+    "            [\"2IWI\"],\n",
+    "            [\"7RPZ\"],\n",
+    "            [\"3TJN\"]\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": "4c61bac4-4f2e-4f4a-aa1f-30dca209747c",
+   "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
+}

app.py

@@ -1,4 +1,11 @@
 import gradio as gr
+import requests
+from Bio.PDB import PDBParser
+import numpy as np
+import os
+from gradio_molecule3d import Molecule3D
+
+
 from model_loader import load_model
 
 import torch
@@ -7,8 +14,6 @@ import torch.nn.functional as F
 from torch.utils.data import DataLoader
 
 import re
-import numpy as np
-import os
 import pandas as pd
 import copy
 
@@ -20,18 +25,6 @@ from datasets import Dataset
 
 from scipy.special import expit
 
-import requests
-
-from gradio_molecule3d import Molecule3D
-
-# Biopython imports
-from Bio.PDB import PDBParser, Select, PDBIO
-from Bio.PDB.DSSP import DSSP
-from Bio.PDB import PDBList
-
-from matplotlib import cm  # For color mapping
-from matplotlib.colors import Normalize
-
 # Load model and move to device
 checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
 max_length = 1500
@@ -40,23 +33,26 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)
 model.eval()
 
-# Function to fetch a PDB file
+def normalize_scores(scores):
+    min_score = np.min(scores)
+    max_score = np.max(scores)
+    return (scores - min_score) / (max_score - min_score) if max_score > min_score else scores
+    
+def read_mol(pdb_path):
+    """Read PDB file and return its content as a string"""
+    with open(pdb_path, 'r') as f:
+        return f.read()
+
 def fetch_pdb(pdb_id):
     pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'
-    pdb_path = f'pdb_files/{pdb_id}.pdb'
-    os.makedirs('pdb_files', exist_ok=True)
+    pdb_path = f'{pdb_id}.pdb'
     response = requests.get(pdb_url)
     if response.status_code == 200:
         with open(pdb_path, 'wb') as f:
             f.write(response.content)
         return pdb_path
-    return None
-
-
-def normalize_scores(scores):
-    min_score = np.min(scores)
-    max_score = np.max(scores)
-    return (scores - min_score) / (max_score - min_score) if max_score > min_score else scores
+    else:
+        return None
 
 def process_pdb(pdb_id, segment):
     pdb_path = fetch_pdb(pdb_id)
@@ -65,7 +61,11 @@ def process_pdb(pdb_id, segment):
     
     parser = PDBParser(QUIET=1)
     structure = parser.get_structure('protein', pdb_path)
-    chain = structure[0][segment]
+    
+    try:
+        chain = structure[0][segment]
+    except KeyError:
+        return "Invalid Chain ID", None, None
     
     # Comprehensive amino acid mapping
     aa_dict = {
@@ -77,11 +77,10 @@ def process_pdb(pdb_id, segment):
     }
     
     # Exclude non-amino acid residues
-    sequence = "".join(
-        aa_dict[residue.get_resname().strip()] 
-        for residue in chain 
+    sequence = [
+        residue for residue in chain 
         if residue.get_resname().strip() in aa_dict
-    )
+    ]
     
     # Prepare input for model prediction
     input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
@@ -91,54 +90,166 @@ def process_pdb(pdb_id, segment):
     # Calculate scores and normalize them
     scores = expit(outputs[:, 1] - outputs[:, 0])
     normalized_scores = normalize_scores(scores)
+
+    result_str = "\n".join(
+        f"{aa_dict[res.get_resname()]} {res.id[1]} {score:.2f}" 
+        for res, score in zip(sequence, normalized_scores)
+    )
     
-    # Prepare the result string, including only amino acid residues
-    result_str = "\n".join([
-        f"{res.get_resname()} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}" 
-        for i, res in enumerate(chain) if res.get_resname().strip() in aa_dict
-    ])
-    
-    # Save predictions to file
-    with open(f"{pdb_id}_predictions.txt", "w") as f:
+    # Save the predictions to a file
+    prediction_file = f"{pdb_id}_predictions.txt"
+    with open(prediction_file, "w") as f:
         f.write(result_str)
     
-    return result_str, pdb_path, f"{pdb_id}_predictions.txt"
+    return result_str, molecule(pdb_path, random_scores, segment), prediction_file
+
+def molecule(input_pdb, scores=None, segment='A'):
+    mol = read_mol(input_pdb)  # Read PDB file content
+    
+    # Prepare high-scoring residues script if scores are provided
+    high_score_script = ""
+    if scores is not None:
+        high_score_script = """
+        // Reset all styles first
+        viewer.getModel(0).setStyle({}, {});
+        
+        // Show only the selected chain
+        viewer.getModel(0).setStyle(
+            {"chain": "%s"}, 
+            { cartoon: {colorscheme:"whiteCarbon"} }
+        );
+        
+        // Highlight high-scoring residues only for the selected chain
+        let highScoreResidues = [%s];
+        viewer.getModel(0).setStyle(
+            {"chain": "%s", "resi": highScoreResidues}, 
+            {"stick": {"color": "red"}}
+        );
 
-reps = [{"model": 0, "style": "cartoon", "color": "spectrum"}]
+        // Highlight high-scoring residues only for the selected chain
+        let highScoreResidues2 = [%s];
+        viewer.getModel(0).setStyle(
+            {"chain": "%s", "resi": highScoreResidues2}, 
+            {"stick": {"color": "orange"}}
+        );
+        """ % (segment, 
+               ", ".join(str(i+1) for i, score in enumerate(scores) if score > 0.8),
+               segment,
+              ", ".join(str(i+1) for i, score in enumerate(scores) if (score > 0.5) and (score < 0.8)),
+               segment)
+    
+    html_content = f"""
+    <!DOCTYPE html>
+    <html>
+    <head>    
+        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+        <style>
+        .mol-container {{
+            width: 100%;
+            height: 700px;
+            position: relative;
+        }}
+        </style>
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js"></script>
+        <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+    </head>
+    <body>
+        <div id="container" class="mol-container"></div>
+        <script>
+            let pdb = `{mol}`; // Use template literal to properly escape PDB content
+            $(document).ready(function () {{
+                let element = $("#container");
+                let config = {{ backgroundColor: "white" }};
+                let viewer = $3Dmol.createViewer(element, config);
+                viewer.addModel(pdb, "pdb");
+                
+                // Reset all styles and show only selected chain
+                viewer.getModel(0).setStyle(
+                    {{"chain": "{segment}"}}, 
+                    {{ cartoon: {{ colorscheme:"whiteCarbon" }} }}
+                );
+                
+                {high_score_script}
+                
+                // Add hover functionality
+                viewer.setHoverable(
+                    {{}}, 
+                    true, 
+                    function(atom, viewer, event, container) {{
+                        if (!atom.label) {{
+                            atom.label = viewer.addLabel(
+                                atom.resn + ":" + atom.atom, 
+                                {{
+                                    position: atom, 
+                                    backgroundColor: 'mintcream', 
+                                    fontColor: 'black',
+                                    fontSize: 12,
+                                    padding: 2
+                                }}
+                            );
+                        }}
+                    }},
+                    function(atom, viewer) {{
+                        if (atom.label) {{
+                            viewer.removeLabel(atom.label);
+                            delete atom.label;
+                        }}
+                    }}
+                );
+                
+                viewer.zoomTo();
+                viewer.render();
+                viewer.zoom(0.8, 2000);
+            }});
+        </script>
+    </body>
+    </html>
+    """
+    
+    # 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>'
+
+reps =    [
+        {
+          "model": 0,
+          "style": "cartoon",
+          "color": "whiteCarbon",
+          "residue_range": "",
+          "around": 0,
+          "byres": False,
+        }
+    ]
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# Protein Binding Site Prediction")
+    gr.Markdown("# Protein Binding Site Prediction (Random Scores)")
+    with gr.Row():
+        pdb_input = gr.Textbox(value="2IWI", label="PDB ID", placeholder="Enter PDB ID here...")
+        visualize_btn = gr.Button("Visualize Structure")
+
+    molecule_output2 = Molecule3D(label="Protein Structure", reps=reps)
 
     with gr.Row():
-        pdb_input = gr.Textbox(value="2IWI",
-                label="PDB ID",
-                placeholder="Enter PDB ID here...")
-        segment_input = gr.Textbox(value="A",
-                label="Chain ID (Segment)",
-                placeholder="Enter Chain ID here...")
-        visualize_btn = gr.Button("Visualize Sructure")
-        prediction_btn = gr.Button("Predict Ligand Binding Site")
-
-    molecule_output = Molecule3D(label="Protein Structure", reps=reps)
+        pdb_input = gr.Textbox(value="2IWI", label="PDB ID", placeholder="Enter PDB ID here...")
+        segment_input = gr.Textbox(value="A", label="Chain ID", placeholder="Enter Chain ID here...")
+        prediction_btn = gr.Button("Predict Random Binding Site Scores")
+
+    molecule_output = gr.HTML(label="Protein Structure")
     predictions_output = gr.Textbox(label="Binding Site Predictions")
     download_output = gr.File(label="Download Predictions")
-
-    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output)
-    prediction_btn.click(
-        process_pdb, 
-        inputs=[pdb_input, segment_input], 
-        outputs=[predictions_output, molecule_output, download_output]
-    )
-
+    
+    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)
+    
+    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])
+    
     gr.Markdown("## Examples")
     gr.Examples(
         examples=[
-            ["2IWI"],
-            ["7RPZ"],
-            ["3TJN"]
+            ["2IWI", "A"],
+            ["7RPZ", "B"],
+            ["3TJN", "C"]
         ],
-        inputs=[pdb_input, segment_input], 
+        inputs=[pdb_input, segment_input],
         outputs=[predictions_output, molecule_output, download_output]
     )
 

model_loader.ipynb

@@ -1,871 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 38,
-   "id": "14ff5741-629c-445a-a8a9-b3d9db1f3ddb",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "import torch.nn as nn\n",
-    "import torch.nn.functional as F\n",
-    "from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss\n",
-    "from torch.utils.data import DataLoader\n",
-    "\n",
-    "import re\n",
-    "import numpy as np\n",
-    "import os\n",
-    "import pandas as pd\n",
-    "import copy\n",
-    "\n",
-    "import transformers, datasets\n",
-    "from transformers.modeling_outputs import TokenClassifierOutput\n",
-    "from transformers.models.t5.modeling_t5 import T5Config, T5PreTrainedModel, T5Stack\n",
-    "from transformers.utils.model_parallel_utils import assert_device_map, get_device_map\n",
-    "from transformers import T5EncoderModel, T5Tokenizer\n",
-    "from transformers.models.esm.modeling_esm import EsmPreTrainedModel, EsmModel\n",
-    "from transformers import AutoTokenizer\n",
-    "from transformers import TrainingArguments, Trainer, set_seed\n",
-    "from transformers import DataCollatorForTokenClassification\n",
-    "\n",
-    "from dataclasses import dataclass\n",
-    "from typing import Dict, List, Optional, Tuple, Union\n",
-    "\n",
-    "# for custom DataCollator\n",
-    "from transformers.data.data_collator import DataCollatorMixin\n",
-    "from transformers.tokenization_utils_base import PreTrainedTokenizerBase\n",
-    "from transformers.utils import PaddingStrategy\n",
-    "\n",
-    "from datasets import Dataset\n",
-    "\n",
-    "from scipy.special import expit\n",
-    "\n",
-    "import peft\n",
-    "from peft import get_peft_config, PeftModel, PeftConfig, inject_adapter_in_model, LoraConfig"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "5ec16a71-ed5d-46a6-98b2-55bc5d0fbe07",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "cnn_head=True #False set True for Rostlab/prot_t5_xl_half_uniref50-enc\n",
-    "ffn_head=False #False\n",
-    "transformer_head=False\n",
-    "custom_lora=True #False #only true for Rostlab/prot_t5_xl_half_uniref50-enc"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "id": "cc7151ca-0daf-4e75-a865-ab52f9b28f2e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "class ClassConfig:\n",
-    "    def __init__(self, dropout=0.2, num_labels=3):\n",
-    "        self.dropout_rate = dropout\n",
-    "        self.num_labels = num_labels\n",
-    "\n",
-    "class T5EncoderForTokenClassification(T5PreTrainedModel):\n",
-    "\n",
-    "    def __init__(self, config: T5Config, class_config: ClassConfig):\n",
-    "        super().__init__(config)\n",
-    "        self.num_labels = class_config.num_labels\n",
-    "        self.config = config\n",
-    "\n",
-    "        self.shared = nn.Embedding(config.vocab_size, config.d_model)\n",
-    "\n",
-    "        encoder_config = copy.deepcopy(config)\n",
-    "        encoder_config.use_cache = False\n",
-    "        encoder_config.is_encoder_decoder = False\n",
-    "        self.encoder = T5Stack(encoder_config, self.shared)\n",
-    "\n",
-    "        self.dropout = nn.Dropout(class_config.dropout_rate)\n",
-    "\n",
-    "        # Initialize different heads based on class_config\n",
-    "        if cnn_head:\n",
-    "            self.cnn = nn.Conv1d(config.hidden_size, 512, kernel_size=3, padding=1)\n",
-    "            self.classifier = nn.Linear(512, class_config.num_labels)\n",
-    "        elif ffn_head:\n",
-    "            # Multi-layer feed-forward network (FFN) head\n",
-    "            self.ffn = nn.Sequential(\n",
-    "                nn.Linear(config.hidden_size, 512),\n",
-    "                nn.ReLU(),\n",
-    "                nn.Linear(512, 256),\n",
-    "                nn.ReLU(),\n",
-    "                nn.Linear(256, class_config.num_labels)\n",
-    "            )\n",
-    "        elif transformer_head:\n",
-    "            # Transformer layer head\n",
-    "            encoder_layer = nn.TransformerEncoderLayer(d_model=config.hidden_size, nhead=8)\n",
-    "            self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)\n",
-    "            self.classifier = nn.Linear(config.hidden_size, class_config.num_labels)\n",
-    "        else:\n",
-    "            # Default classification head\n",
-    "            self.classifier = nn.Linear(config.hidden_size, class_config.num_labels)\n",
-    "        \n",
-    "        self.post_init()\n",
-    "\n",
-    "        # Model parallel\n",
-    "        self.model_parallel = False\n",
-    "        self.device_map = None\n",
-    "\n",
-    "    def parallelize(self, device_map=None):\n",
-    "        self.device_map = (\n",
-    "            get_device_map(len(self.encoder.block), range(torch.cuda.device_count()))\n",
-    "            if device_map is None\n",
-    "            else device_map\n",
-    "        )\n",
-    "        assert_device_map(self.device_map, len(self.encoder.block))\n",
-    "        self.encoder.parallelize(self.device_map)\n",
-    "        self.classifier = self.classifier.to(self.encoder.first_device)\n",
-    "        self.model_parallel = True\n",
-    "\n",
-    "    def deparallelize(self):\n",
-    "        self.encoder.deparallelize()\n",
-    "        self.encoder = self.encoder.to(\"cpu\")\n",
-    "        self.model_parallel = False\n",
-    "        self.device_map = None\n",
-    "        torch.cuda.empty_cache()\n",
-    "\n",
-    "    def get_input_embeddings(self):\n",
-    "        return self.shared\n",
-    "\n",
-    "    def set_input_embeddings(self, new_embeddings):\n",
-    "        self.shared = new_embeddings\n",
-    "        self.encoder.set_input_embeddings(new_embeddings)\n",
-    "\n",
-    "    def get_encoder(self):\n",
-    "        return self.encoder\n",
-    "\n",
-    "    def _prune_heads(self, heads_to_prune):\n",
-    "        for layer, heads in heads_to_prune.items():\n",
-    "            self.encoder.layer[layer].attention.prune_heads(heads)\n",
-    "\n",
-    "    def forward(\n",
-    "        self,\n",
-    "        input_ids=None,\n",
-    "        attention_mask=None,\n",
-    "        head_mask=None,\n",
-    "        inputs_embeds=None,\n",
-    "        labels=None,\n",
-    "        output_attentions=None,\n",
-    "        output_hidden_states=None,\n",
-    "        return_dict=None,\n",
-    "    ):\n",
-    "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
-    "\n",
-    "        outputs = self.encoder(\n",
-    "            input_ids=input_ids,\n",
-    "            attention_mask=attention_mask,\n",
-    "            inputs_embeds=inputs_embeds,\n",
-    "            head_mask=head_mask,\n",
-    "            output_attentions=output_attentions,\n",
-    "            output_hidden_states=output_hidden_states,\n",
-    "            return_dict=return_dict,\n",
-    "        )\n",
-    "\n",
-    "        sequence_output = outputs[0]\n",
-    "        sequence_output = self.dropout(sequence_output)\n",
-    "\n",
-    "        # Forward pass through the selected head\n",
-    "        if cnn_head:\n",
-    "            # CNN head\n",
-    "            sequence_output = sequence_output.permute(0, 2, 1)  # Prepare shape for CNN\n",
-    "            cnn_output = self.cnn(sequence_output)\n",
-    "            cnn_output = F.relu(cnn_output)\n",
-    "            cnn_output = cnn_output.permute(0, 2, 1)  # Shape back for classifier\n",
-    "            logits = self.classifier(cnn_output)\n",
-    "        elif ffn_head:\n",
-    "            # FFN head\n",
-    "            logits = self.ffn(sequence_output)\n",
-    "        elif transformer_head:\n",
-    "            # Transformer head\n",
-    "            transformer_output = self.transformer_encoder(sequence_output)\n",
-    "            logits = self.classifier(transformer_output)\n",
-    "        else:\n",
-    "            # Default classification head\n",
-    "            logits = self.classifier(sequence_output)\n",
-    "\n",
-    "        loss = None\n",
-    "        if labels is not None:\n",
-    "            loss_fct = CrossEntropyLoss()\n",
-    "            active_loss = attention_mask.view(-1) == 1\n",
-    "            active_logits = logits.view(-1, self.num_labels)\n",
-    "            active_labels = torch.where(\n",
-    "                active_loss, labels.view(-1), torch.tensor(-100).type_as(labels)\n",
-    "            )\n",
-    "            valid_logits = active_logits[active_labels != -100]\n",
-    "            valid_labels = active_labels[active_labels != -100]\n",
-    "            valid_labels = valid_labels.to(valid_logits.device)\n",
-    "            valid_labels = valid_labels.long()\n",
-    "            loss = loss_fct(valid_logits, valid_labels)\n",
-    "\n",
-    "        if not return_dict:\n",
-    "            output = (logits,) + outputs[2:]\n",
-    "            return ((loss,) + output) if loss is not None else output\n",
-    "\n",
-    "        return TokenClassifierOutput(\n",
-    "            loss=loss,\n",
-    "            logits=logits,\n",
-    "            hidden_states=outputs.hidden_states,\n",
-    "            attentions=outputs.attentions,\n",
-    "        )"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "id": "e5e751ba-f4d3-4a28-bea0-82633f1dabb4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Modifies an existing transformer and introduce the LoRA layers\n",
-    "\n",
-    "class CustomLoRAConfig:\n",
-    "    def __init__(self):\n",
-    "        self.lora_rank = 4\n",
-    "        self.lora_init_scale = 0.01\n",
-    "        self.lora_modules = \".*SelfAttention|.*EncDecAttention\"\n",
-    "        self.lora_layers = \"q|k|v|o\"\n",
-    "        self.trainable_param_names = \".*layer_norm.*|.*lora_[ab].*\"\n",
-    "        self.lora_scaling_rank = 1\n",
-    "        # lora_modules and lora_layers are speicified with regular expressions\n",
-    "        # see https://www.w3schools.com/python/python_regex.asp for reference\n",
-    "        \n",
-    "class LoRALinear(nn.Module):\n",
-    "    def __init__(self, linear_layer, rank, scaling_rank, init_scale):\n",
-    "        super().__init__()\n",
-    "        self.in_features = linear_layer.in_features\n",
-    "        self.out_features = linear_layer.out_features\n",
-    "        self.rank = rank\n",
-    "        self.scaling_rank = scaling_rank\n",
-    "        self.weight = linear_layer.weight\n",
-    "        self.bias = linear_layer.bias\n",
-    "        if self.rank > 0:\n",
-    "            self.lora_a = nn.Parameter(torch.randn(rank, linear_layer.in_features) * init_scale)\n",
-    "            if init_scale < 0:\n",
-    "                self.lora_b = nn.Parameter(torch.randn(linear_layer.out_features, rank) * init_scale)\n",
-    "            else:\n",
-    "                self.lora_b = nn.Parameter(torch.zeros(linear_layer.out_features, rank))\n",
-    "        if self.scaling_rank:\n",
-    "            self.multi_lora_a = nn.Parameter(\n",
-    "                torch.ones(self.scaling_rank, linear_layer.in_features)\n",
-    "                + torch.randn(self.scaling_rank, linear_layer.in_features) * init_scale\n",
-    "            )\n",
-    "            if init_scale < 0:\n",
-    "                self.multi_lora_b = nn.Parameter(\n",
-    "                    torch.ones(linear_layer.out_features, self.scaling_rank)\n",
-    "                    + torch.randn(linear_layer.out_features, self.scaling_rank) * init_scale\n",
-    "                )\n",
-    "            else:\n",
-    "                self.multi_lora_b = nn.Parameter(torch.ones(linear_layer.out_features, self.scaling_rank))\n",
-    "\n",
-    "    def forward(self, input):\n",
-    "        if self.scaling_rank == 1 and self.rank == 0:\n",
-    "            # parsimonious implementation for ia3 and lora scaling\n",
-    "            if self.multi_lora_a.requires_grad:\n",
-    "                hidden = F.linear((input * self.multi_lora_a.flatten()), self.weight, self.bias)\n",
-    "            else:\n",
-    "                hidden = F.linear(input, self.weight, self.bias)\n",
-    "            if self.multi_lora_b.requires_grad:\n",
-    "                hidden = hidden * self.multi_lora_b.flatten()\n",
-    "            return hidden\n",
-    "        else:\n",
-    "            # general implementation for lora (adding and scaling)\n",
-    "            weight = self.weight\n",
-    "            if self.scaling_rank:\n",
-    "                weight = weight * torch.matmul(self.multi_lora_b, self.multi_lora_a) / self.scaling_rank\n",
-    "            if self.rank:\n",
-    "                weight = weight + torch.matmul(self.lora_b, self.lora_a) / self.rank\n",
-    "            return F.linear(input, weight, self.bias)\n",
-    "\n",
-    "    def extra_repr(self):\n",
-    "        return \"in_features={}, out_features={}, bias={}, rank={}, scaling_rank={}\".format(\n",
-    "            self.in_features, self.out_features, self.bias is not None, self.rank, self.scaling_rank\n",
-    "        )\n",
-    "\n",
-    "\n",
-    "def modify_with_lora(transformer, config):\n",
-    "    for m_name, module in dict(transformer.named_modules()).items():\n",
-    "        if re.fullmatch(config.lora_modules, m_name):\n",
-    "            for c_name, layer in dict(module.named_children()).items():\n",
-    "                if re.fullmatch(config.lora_layers, c_name):\n",
-    "                    assert isinstance(\n",
-    "                        layer, nn.Linear\n",
-    "                    ), f\"LoRA can only be applied to torch.nn.Linear, but {layer} is {type(layer)}.\"\n",
-    "                    setattr(\n",
-    "                        module,\n",
-    "                        c_name,\n",
-    "                        LoRALinear(layer, config.lora_rank, config.lora_scaling_rank, config.lora_init_scale),\n",
-    "                    )\n",
-    "    return transformer\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "id": "43a56311-3279-466a-bc95-590381f1b13c",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def load_T5_model_classification(checkpoint, num_labels, half_precision, full = False, deepspeed=True):\n",
-    "    # Load model and tokenizer\n",
-    "\n",
-    "    if \"ankh\" in checkpoint :\n",
-    "        model = T5EncoderModel.from_pretrained(checkpoint)\n",
-    "        tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
-    "\n",
-    "    elif \"prot_t5\" in checkpoint:\n",
-    "        # possible to load the half precision model (thanks to @pawel-rezo for pointing that out)\n",
-    "        if half_precision and deepspeed:\n",
-    "            #tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc', do_lower_case=False)\n",
-    "            #model = T5EncoderModel.from_pretrained(\"Rostlab/prot_t5_xl_half_uniref50-enc\", torch_dtype=torch.float16)#.to(torch.device('cuda')\n",
-    "            tokenizer = T5Tokenizer.from_pretrained(checkpoint, do_lower_case=False)\n",
-    "            model = T5EncoderModel.from_pretrained(checkpoint, torch_dtype=torch.float16).to(torch.device('cuda'))\n",
-    "        else:\n",
-    "            model = T5EncoderModel.from_pretrained(checkpoint)\n",
-    "            tokenizer = T5Tokenizer.from_pretrained(checkpoint)\n",
-    "                \n",
-    "    elif \"ProstT5\" in checkpoint:\n",
-    "        if half_precision and deepspeed: \n",
-    "            tokenizer = T5Tokenizer.from_pretrained(checkpoint, do_lower_case=False)\n",
-    "            model = T5EncoderModel.from_pretrained(checkpoint, torch_dtype=torch.float16).to(torch.device('cuda'))\n",
-    "        else:\n",
-    "            model = T5EncoderModel.from_pretrained(checkpoint)\n",
-    "            tokenizer = T5Tokenizer.from_pretrained(checkpoint) \n",
-    "    \n",
-    "    # Create new Classifier model with PT5 dimensions\n",
-    "    class_config=ClassConfig(num_labels=num_labels)\n",
-    "    class_model=T5EncoderForTokenClassification(model.config,class_config)\n",
-    "    \n",
-    "    # Set encoder and embedding weights to checkpoint weights\n",
-    "    class_model.shared=model.shared\n",
-    "    class_model.encoder=model.encoder    \n",
-    "    \n",
-    "    # Delete the checkpoint model\n",
-    "    model=class_model\n",
-    "    del class_model\n",
-    "    \n",
-    "    if full == True:\n",
-    "        return model, tokenizer \n",
-    "    \n",
-    "    # Print number of trainable parameters\n",
-    "    model_parameters = filter(lambda p: p.requires_grad, model.parameters())\n",
-    "    params = sum([np.prod(p.size()) for p in model_parameters])\n",
-    "    print(\"T5_Classfier\\nTrainable Parameter: \"+ str(params))    \n",
-    "\n",
-    "    if custom_lora:\n",
-    "        #the linear CustomLoRAConfig allows better quality predictions, but more memory is needed\n",
-    "        # Add model modification lora\n",
-    "        config = CustomLoRAConfig()\n",
-    "        \n",
-    "        # Add LoRA layers\n",
-    "        model = modify_with_lora(model, config)\n",
-    "        \n",
-    "        # Freeze Embeddings and Encoder (except LoRA)\n",
-    "        for (param_name, param) in model.shared.named_parameters():\n",
-    "                    param.requires_grad = False\n",
-    "        for (param_name, param) in model.encoder.named_parameters():\n",
-    "                    param.requires_grad = False       \n",
-    "    \n",
-    "        for (param_name, param) in model.named_parameters():\n",
-    "                if re.fullmatch(config.trainable_param_names, param_name):\n",
-    "                    param.requires_grad = True\n",
-    "\n",
-    "    else:\n",
-    "        # lora modification\n",
-    "        peft_config = LoraConfig(\n",
-    "            r=4, lora_alpha=1, bias=\"all\", target_modules=[\"q\",\"k\",\"v\",\"o\"]\n",
-    "        )\n",
-    "        \n",
-    "        model = inject_adapter_in_model(peft_config, model)\n",
-    "        \n",
-    "        # Unfreeze the prediction head\n",
-    "        for (param_name, param) in model.classifier.named_parameters():\n",
-    "                    param.requires_grad = True  \n",
-    "\n",
-    "    # Print trainable Parameter          \n",
-    "    model_parameters = filter(lambda p: p.requires_grad, model.parameters())\n",
-    "    params = sum([np.prod(p.size()) for p in model_parameters])\n",
-    "    print(\"T5_LoRA_Classfier\\nTrainable Parameter: \"+ str(params) + \"\\n\")\n",
-    "    \n",
-    "    return model, tokenizer"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "id": "7ba720bc-a003-4984-a965-cb2f42344e85",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "class EsmForTokenClassificationCustom(EsmPreTrainedModel):\n",
-    "    _keys_to_ignore_on_load_unexpected = [r\"pooler\"]\n",
-    "    _keys_to_ignore_on_load_missing = [r\"position_ids\", r\"cnn\", r\"ffn\", r\"transformer\"]\n",
-    "\n",
-    "    def __init__(self, config):\n",
-    "        super().__init__(config)\n",
-    "        self.num_labels = config.num_labels\n",
-    "        self.esm = EsmModel(config, add_pooling_layer=False)\n",
-    "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
-    "\n",
-    "        if cnn_head:\n",
-    "            self.cnn = nn.Conv1d(config.hidden_size, 512, kernel_size=3, padding=1)\n",
-    "            self.classifier = nn.Linear(512, config.num_labels)\n",
-    "        elif ffn_head:\n",
-    "            # Multi-layer feed-forward network (FFN) as an alternative head\n",
-    "            self.ffn = nn.Sequential(\n",
-    "                nn.Linear(config.hidden_size, 512),\n",
-    "                nn.ReLU(),\n",
-    "                nn.Linear(512, 256),\n",
-    "                nn.ReLU(),\n",
-    "                nn.Linear(256, config.num_labels)\n",
-    "            )\n",
-    "        elif transformer_head:\n",
-    "            # Transformer layer as an alternative head\n",
-    "            encoder_layer = nn.TransformerEncoderLayer(d_model=config.hidden_size, nhead=8)\n",
-    "            self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)\n",
-    "            self.classifier = nn.Linear(config.hidden_size, config.num_labels)\n",
-    "        else:\n",
-    "            # Default classification head\n",
-    "            self.classifier = nn.Linear(config.hidden_size, config.num_labels)\n",
-    "\n",
-    "        self.init_weights()\n",
-    "\n",
-    "    def forward(\n",
-    "        self,\n",
-    "        input_ids: Optional[torch.LongTensor] = None,\n",
-    "        attention_mask: Optional[torch.Tensor] = None,\n",
-    "        position_ids: Optional[torch.LongTensor] = None,\n",
-    "        head_mask: Optional[torch.Tensor] = None,\n",
-    "        inputs_embeds: Optional[torch.FloatTensor] = None,\n",
-    "        labels: Optional[torch.LongTensor] = None,\n",
-    "        output_attentions: Optional[bool] = None,\n",
-    "        output_hidden_states: Optional[bool] = None,\n",
-    "        return_dict: Optional[bool] = None,\n",
-    "    ) -> Union[Tuple, TokenClassifierOutput]:\n",
-    "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
-    "        outputs = self.esm(\n",
-    "            input_ids,\n",
-    "            attention_mask=attention_mask,\n",
-    "            position_ids=position_ids,\n",
-    "            head_mask=head_mask,\n",
-    "            inputs_embeds=inputs_embeds,\n",
-    "            output_attentions=output_attentions,\n",
-    "            output_hidden_states=output_hidden_states,\n",
-    "            return_dict=return_dict,\n",
-    "        )\n",
-    "        \n",
-    "        sequence_output = outputs[0]\n",
-    "        sequence_output = self.dropout(sequence_output)\n",
-    "\n",
-    "        if cnn_head:\n",
-    "            sequence_output = sequence_output.transpose(1, 2)\n",
-    "            sequence_output = self.cnn(sequence_output)\n",
-    "            sequence_output = sequence_output.transpose(1, 2)\n",
-    "            logits = self.classifier(sequence_output)\n",
-    "        elif ffn_head:\n",
-    "            logits = self.ffn(sequence_output)\n",
-    "        elif transformer_head:\n",
-    "            # Apply transformer encoder for the transformer head\n",
-    "            sequence_output = self.transformer_encoder(sequence_output)\n",
-    "            logits = self.classifier(sequence_output)\n",
-    "        else:\n",
-    "            logits = self.classifier(sequence_output)\n",
-    "\n",
-    "        loss = None\n",
-    "        if labels is not None:\n",
-    "            loss_fct = CrossEntropyLoss()\n",
-    "            active_loss = attention_mask.view(-1) == 1\n",
-    "            active_logits = logits.view(-1, self.num_labels)\n",
-    "            active_labels = torch.where(\n",
-    "                active_loss, labels.view(-1), torch.tensor(-100).type_as(labels)\n",
-    "            )\n",
-    "            valid_logits = active_logits[active_labels != -100]\n",
-    "            valid_labels = active_labels[active_labels != -100]\n",
-    "            valid_labels = valid_labels.type(torch.LongTensor).to('cuda:0')\n",
-    "            loss = loss_fct(valid_logits, valid_labels)\n",
-    "\n",
-    "        if not return_dict:\n",
-    "            output = (logits,) + outputs[2:]\n",
-    "            return ((loss,) + output) if loss is not None else output\n",
-    "\n",
-    "        return TokenClassifierOutput(\n",
-    "            loss=loss,\n",
-    "            logits=logits,\n",
-    "            hidden_states=outputs.hidden_states,\n",
-    "            attentions=outputs.attentions,\n",
-    "        )\n",
-    "\n",
-    "    def _init_weights(self, module):\n",
-    "        if isinstance(module, nn.Linear) or isinstance(module, nn.Conv1d):\n",
-    "            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)\n",
-    "            if module.bias is not None:\n",
-    "                module.bias.data.zero_()\n",
-    "\n",
-    "# based on transformers DataCollatorForTokenClassification\n",
-    "@dataclass\n",
-    "class DataCollatorForTokenClassificationESM(DataCollatorMixin):\n",
-    "    \"\"\"\n",
-    "    Data collator that will dynamically pad the inputs received, as well as the labels.\n",
-    "    Args:\n",
-    "        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):\n",
-    "            The tokenizer used for encoding the data.\n",
-    "        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):\n",
-    "            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)\n",
-    "            among:\n",
-    "            - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single\n",
-    "              sequence is provided).\n",
-    "            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum\n",
-    "              acceptable input length for the model if that argument is not provided.\n",
-    "            - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths).\n",
-    "        max_length (`int`, *optional*):\n",
-    "            Maximum length of the returned list and optionally padding length (see above).\n",
-    "        pad_to_multiple_of (`int`, *optional*):\n",
-    "            If set will pad the sequence to a multiple of the provided value.\n",
-    "            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=\n",
-    "            7.5 (Volta).\n",
-    "        label_pad_token_id (`int`, *optional*, defaults to -100):\n",
-    "            The id to use when padding the labels (-100 will be automatically ignore by PyTorch loss functions).\n",
-    "        return_tensors (`str`):\n",
-    "            The type of Tensor to return. Allowable values are \"np\", \"pt\" and \"tf\".\n",
-    "    \"\"\"\n",
-    "\n",
-    "    tokenizer: PreTrainedTokenizerBase\n",
-    "    padding: Union[bool, str, PaddingStrategy] = True\n",
-    "    max_length: Optional[int] = None\n",
-    "    pad_to_multiple_of: Optional[int] = None\n",
-    "    label_pad_token_id: int = -100\n",
-    "    return_tensors: str = \"pt\"\n",
-    "\n",
-    "    def torch_call(self, features):\n",
-    "        import torch\n",
-    "\n",
-    "        label_name = \"label\" if \"label\" in features[0].keys() else \"labels\"\n",
-    "        labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None\n",
-    "\n",
-    "        no_labels_features = [{k: v for k, v in feature.items() if k != label_name} for feature in features]\n",
-    "\n",
-    "        batch = self.tokenizer.pad(\n",
-    "            no_labels_features,\n",
-    "            padding=self.padding,\n",
-    "            max_length=self.max_length,\n",
-    "            pad_to_multiple_of=self.pad_to_multiple_of,\n",
-    "            return_tensors=\"pt\",\n",
-    "        )\n",
-    "\n",
-    "        if labels is None:\n",
-    "            return batch\n",
-    "\n",
-    "        sequence_length = batch[\"input_ids\"].shape[1]\n",
-    "        padding_side = self.tokenizer.padding_side\n",
-    "\n",
-    "        def to_list(tensor_or_iterable):\n",
-    "            if isinstance(tensor_or_iterable, torch.Tensor):\n",
-    "                return tensor_or_iterable.tolist()\n",
-    "            return list(tensor_or_iterable)\n",
-    "\n",
-    "        if padding_side == \"right\":\n",
-    "            batch[label_name] = [\n",
-    "                # to_list(label) + [self.label_pad_token_id] * (sequence_length - len(label)) for label in labels\n",
-    "                # changed to pad the special tokens at the beginning and end of the sequence\n",
-    "                [self.label_pad_token_id] + to_list(label) + [self.label_pad_token_id] * (sequence_length - len(label)-1) for label in labels\n",
-    "            ]\n",
-    "        else:\n",
-    "            batch[label_name] = [\n",
-    "                [self.label_pad_token_id] * (sequence_length - len(label)) + to_list(label) for label in labels\n",
-    "            ]\n",
-    "\n",
-    "        batch[label_name] = torch.tensor(batch[label_name], dtype=torch.float)\n",
-    "        return batch\n",
-    "\n",
-    "def _torch_collate_batch(examples, tokenizer, pad_to_multiple_of: Optional[int] = None):\n",
-    "    \"\"\"Collate `examples` into a batch, using the information in `tokenizer` for padding if necessary.\"\"\"\n",
-    "    import torch\n",
-    "\n",
-    "    # Tensorize if necessary.\n",
-    "    if isinstance(examples[0], (list, tuple, np.ndarray)):\n",
-    "        examples = [torch.tensor(e, dtype=torch.long) for e in examples]\n",
-    "\n",
-    "    length_of_first = examples[0].size(0)\n",
-    "\n",
-    "    # Check if padding is necessary.\n",
-    "\n",
-    "    are_tensors_same_length = all(x.size(0) == length_of_first for x in examples)\n",
-    "    if are_tensors_same_length and (pad_to_multiple_of is None or length_of_first % pad_to_multiple_of == 0):\n",
-    "        return torch.stack(examples, dim=0)\n",
-    "\n",
-    "    # If yes, check if we have a `pad_token`.\n",
-    "    if tokenizer._pad_token is None:\n",
-    "        raise ValueError(\n",
-    "            \"You are attempting to pad samples but the tokenizer you are using\"\n",
-    "            f\" ({tokenizer.__class__.__name__}) does not have a pad token.\"\n",
-    "        )\n",
-    "\n",
-    "    # Creating the full tensor and filling it with our data.\n",
-    "    max_length = max(x.size(0) for x in examples)\n",
-    "    if pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):\n",
-    "        max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of\n",
-    "    result = examples[0].new_full([len(examples), max_length], tokenizer.pad_token_id)\n",
-    "    for i, example in enumerate(examples):\n",
-    "        if tokenizer.padding_side == \"right\":\n",
-    "            result[i, : example.shape[0]] = example\n",
-    "        else:\n",
-    "            result[i, -example.shape[0] :] = example\n",
-    "    return result\n",
-    "\n",
-    "def tolist(x):\n",
-    "    if isinstance(x, list):\n",
-    "        return x\n",
-    "    elif hasattr(x, \"numpy\"):  # Checks for TF tensors without needing the import\n",
-    "        x = x.numpy()\n",
-    "    return x.tolist()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "id": "ea511812-1244-4e51-b63c-b4da7822f0b7",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "#load ESM2 models\n",
-    "def load_esm_model_classification(checkpoint, num_labels, half_precision, full=False, deepspeed=True):\n",
-    "    \n",
-    "    tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
-    "\n",
-    "    \n",
-    "    if half_precision and deepspeed:\n",
-    "        model = EsmForTokenClassificationCustom.from_pretrained(checkpoint, \n",
-    "                                                                num_labels = num_labels, \n",
-    "                                                                ignore_mismatched_sizes=True,\n",
-    "                                                                torch_dtype = torch.float16)\n",
-    "    else:\n",
-    "        model = EsmForTokenClassificationCustom.from_pretrained(checkpoint, \n",
-    "                                                                num_labels = num_labels,\n",
-    "                                                                ignore_mismatched_sizes=True)\n",
-    "        \n",
-    "    if full == True:\n",
-    "        return model, tokenizer \n",
-    "        \n",
-    "    peft_config = LoraConfig(\n",
-    "        r=4, lora_alpha=1, bias=\"all\", target_modules=[\"query\",\"key\",\"value\",\"dense\"]\n",
-    "    )\n",
-    "    \n",
-    "    model = inject_adapter_in_model(peft_config, model)\n",
-    "\n",
-    "    #model.gradient_checkpointing_enable()\n",
-    "    \n",
-    "    # Unfreeze the prediction head\n",
-    "    for (param_name, param) in model.classifier.named_parameters():\n",
-    "                param.requires_grad = True  \n",
-    "    \n",
-    "    return model, tokenizer"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 22,
-   "id": "8941bbbb-57c5-4f3d-89d9-12b2d306e7a1",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "checkpoint='../Pretrained/Rostlab/prot_t5_xl_uniref50'\n",
-    "best_model_path='../refined_models/ChallengeFinetuning/Rostlab/prot_t5_xl_uniref50/manual_checkpoint/cpt.pth'\n",
-    "full=False\n",
-    "deepspeed=False\n",
-    "mixed=False \n",
-    "num_labels=2"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4f007331-34d4-4c1d-9311-e91db23d9ed5",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "/home/frohlkin/Projects/PLM/Publication/hf_webpage/pretrained"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 24,
-   "id": "18d4ad06-b195-4cc6-a3c8-fa3e761838dc",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "../Pretrained/Rostlab/prot_t5_xl_uniref50 2 False False False\n",
-      "T5_Classfier\n",
-      "Trainable Parameter: 1209716226\n",
-      "T5_LoRA_Classfier\n",
-      "Trainable Parameter: 4082178\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "<All keys matched successfully>"
-      ]
-     },
-     "execution_count": 24,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "print(checkpoint, num_labels, mixed, full, deepspeed)\n",
-    "        \n",
-    "# Determine model type and load accordingly\n",
-    "if \"esm\" in checkpoint:\n",
-    "    model, tokenizer = load_esm_model_classification(checkpoint, num_labels, mixed, full, deepspeed)\n",
-    "else:\n",
-    "    model, tokenizer = load_T5_model_classification(checkpoint, num_labels, mixed, full, deepspeed)\n",
-    "\n",
-    "# Load the best model state\n",
-    "state_dict = torch.load(best_model_path, weights_only=True)\n",
-    "model.load_state_dict(state_dict)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 30,
-   "id": "4e215923-dfe2-4426-aedf-5cb81f7f0db2",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "test_one_letter_sequence='AWYAAK'\n",
-    "max_length=1500"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 40,
-   "id": "7174ea02-ed51-46f5-84c0-6bcd760670d4",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "(7,)"
-      ]
-     },
-     "execution_count": 40,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "def create_dataset(tokenizer,seqs,labels,checkpoint):\n",
-    "    \n",
-    "    tokenized = tokenizer(seqs, max_length=max_length, padding=False, truncation=True)\n",
-    "    dataset = Dataset.from_dict(tokenized)\n",
-    "    \n",
-    "    if (\"esm\" in checkpoint) or (\"ProstT5\" in checkpoint):\n",
-    "        labels = [l[:max_length-2] for l in labels] \n",
-    "    else:\n",
-    "        labels = [l[:max_length-1] for l in labels] \n",
-    "        \n",
-    "    dataset = dataset.add_column(\"labels\", labels)\n",
-    "     \n",
-    "    return dataset\n",
-    "    \n",
-    "def convert_predictions(input_logits):\n",
-    "    all_probs = []\n",
-    "    for logits in input_logits:\n",
-    "        logits = logits.reshape(-1, 2)\n",
-    "\n",
-    "        # Mask out irrelevant regions\n",
-    "        # Compute probabilities for class 1\n",
-    "        probabilities_class1 = expit(logits[:, 1] - logits[:, 0])\n",
-    "        \n",
-    "        all_probs.append(probabilities_class1)\n",
-    "        \n",
-    "    return np.concatenate(all_probs)\n",
-    "    \n",
-    "    \n",
-    "dummy_labels=[np.zeros(len(test_one_letter_sequence))]\n",
-    "# Replace uncommon amino acids with \"X\"\n",
-    "test_one_letter_sequence = test_one_letter_sequence.replace(\"O\", \"X\").replace(\"B\", \"X\").replace(\"U\", \"X\").replace(\"Z\", \"X\").replace(\"J\", \"X\")\n",
-    "\n",
-    "# Add spaces between each amino acid for ProtT5 and ProstT5 models\n",
-    "if \"Rostlab\" in checkpoint:\n",
-    "    test_one_letter_sequence = \" \".join(test_one_letter_sequence)\n",
-    "\n",
-    "# Add <AA2fold> for ProstT5 model input format\n",
-    "if \"ProstT5\" in checkpoint:\n",
-    "    test_one_letter_sequence = \"<AA2fold> \" + test_one_letter_sequence\n",
-    "    \n",
-    "test_dataset=create_dataset(tokenizer,[test_one_letter_sequence],dummy_labels,checkpoint)\n",
-    "\n",
-    "if (\"esm\" in checkpoint) or (\"ProstT5\" in checkpoint):\n",
-    "    data_collator = DataCollatorForTokenClassificationESM(tokenizer)\n",
-    "else:\n",
-    "    data_collator = DataCollatorForTokenClassification(tokenizer)\n",
-    "\n",
-    "test_loader = DataLoader(test_dataset, batch_size=1, collate_fn=data_collator)\n",
-    "\n",
-    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
-    "model.to(device)\n",
-    "for batch in test_loader:\n",
-    "    input_ids = batch['input_ids'].to(device)\n",
-    "    attention_mask = batch['attention_mask'].to(device)\n",
-    "    labels = batch['labels']  # Ensure to get labels from batch\n",
-    "\n",
-    "    outputs = model(input_ids, attention_mask=attention_mask)\n",
-    "    logits = outputs.logits.detach().cpu().numpy()\n",
-    "\n",
-    "logits=convert_predictions(logits)\n",
-    "logits.shape\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",
-    "normalized_scores = normalize_scores(logits)\n",
-    "\n",
-    "normalized_scores.shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "58b5ae4d-9e8e-4d07-ab46-76d23cc29016",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python [conda env:LLM] *",
-   "language": "python",
-   "name": "conda-env-LLM-py"
-  },
-  "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.2"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

requirements.txt

@@ -10,5 +10,4 @@ sentencepiece
 huggingface_hub>=0.15.0
 requests
 gradio_molecule3d
-biopython>=1.81
-matplotlib
+biopython>=1.81

test.ipynb

@@ -0,0 +1,846 @@
+{
+ "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

@@ -0,0 +1,1193 @@
+{
+ "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": 11,
+   "id": "d62be1b5-762e-4b69-aed4-e4ba2a44482f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on local URL:  http://127.0.0.1:7867\n",
+      "\n",
+      "To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7867/\" 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": 11,
+     "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.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": "30f35243-852f-4771-9a4b-5cdd198552b5",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5eca6754-4aa1-463f-881a-25d2a0d6bb5b",
+   "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",
+    "    # 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",
+    "    # 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",
+    "    result_str = \"\\n\".join(\n",
+    "        f\"{aa_dict[res.get_resname()]} {res.id[1]} {score:.2f}\" \n",
+    "        for res, score in zip(sequence, normalized_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.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": "95046d1c-ec7c-4e3e-8a98-1802cb09a25b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a37cbe6f-d57f-41e5-8ae1-38258da39d47",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import gradio as gr\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 numpy as np\n",
+    "import os\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",
+    "import requests\n",
+    "\n",
+    "from gradio_molecule3d import Molecule3D\n",
+    "\n",
+    "# Biopython imports\n",
+    "from Bio.PDB import PDBParser, Select, PDBIO\n",
+    "from Bio.PDB.DSSP import DSSP\n",
+    "from Bio.PDB import PDBList\n",
+    "\n",
+    "from matplotlib import cm  # For color mapping\n",
+    "from matplotlib.colors import Normalize\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",
+    "# Function to fetch a PDB file\n",
+    "def fetch_pdb(pdb_id):\n",
+    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
+    "    pdb_path = f'pdb_files/{pdb_id}.pdb'\n",
+    "    os.makedirs('pdb_files', exist_ok=True)\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",
+    "    return None\n",
+    "\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 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",
+    "    chain = structure[0][segment]\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 = \"\".join(\n",
+    "        aa_dict[residue.get_resname().strip()] \n",
+    "        for residue in chain \n",
+    "        if residue.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",
+    "    # Prepare the result string, including only amino acid residues\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 predictions to file\n",
+    "    with open(f\"{pdb_id}_predictions.txt\", \"w\") as f:\n",
+    "        f.write(result_str)\n",
+    "    \n",
+    "    return result_str, pdb_path, f\"{pdb_id}_predictions.txt\"\n",
+    "\n",
+    "reps = [{\"model\": 0, \"style\": \"cartoon\", \"color\": \"spectrum\"}]\n",
+    "\n",
+    "# Gradio UI\n",
+    "with gr.Blocks() as demo:\n",
+    "    gr.Markdown(\"# Protein Binding Site Prediction\")\n",
+    "\n",
+    "    with gr.Row():\n",
+    "        pdb_input = gr.Textbox(value=\"2IWI\",\n",
+    "                label=\"PDB ID\",\n",
+    "                placeholder=\"Enter PDB ID here...\")\n",
+    "        segment_input = gr.Textbox(value=\"A\",\n",
+    "                label=\"Chain ID (Segment)\",\n",
+    "                placeholder=\"Enter Chain ID here...\")\n",
+    "        visualize_btn = gr.Button(\"Visualize Sructure\")\n",
+    "        prediction_btn = gr.Button(\"Predict Ligand Binding Site\")\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(\n",
+    "        process_pdb, \n",
+    "        inputs=[pdb_input, segment_input], \n",
+    "        outputs=[predictions_output, molecule_output, download_output]\n",
+    "    )\n",
+    "\n",
+    "    gr.Markdown(\"## Examples\")\n",
+    "    gr.Examples(\n",
+    "        examples=[\n",
+    "            [\"2IWI\"],\n",
+    "            [\"7RPZ\"],\n",
+    "            [\"3TJN\"]\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": "4c61bac4-4f2e-4f4a-aa1f-30dca209747c",
+   "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
+}