Update

.ipynb_checkpoints/app-checkpoint.py

@@ -82,6 +82,10 @@ def process_pdb(pdb_id, segment):
         for residue in chain 
         if residue.get_resname().strip() in aa_dict
     )
+    sequence2 = [
+        (res.id[1], res) for res in chain
+        if res.get_resname().strip() in aa_dict
+    ]
     
     # Prepare input for model prediction
     input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
@@ -92,6 +96,9 @@ def process_pdb(pdb_id, segment):
     scores = expit(outputs[:, 1] - outputs[:, 0])
     normalized_scores = normalize_scores(scores)
 
+    # Zip residues with scores to track the residue ID and score
+    residue_scores = [(resi, score) for (resi, _), score in zip(sequence2, normalized_scores)]
+    
     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
@@ -102,14 +109,18 @@ def process_pdb(pdb_id, segment):
     with open(prediction_file, "w") as f:
         f.write(result_str)
     
-    return result_str, molecule(pdb_path, normalized_scores, segment), prediction_file
+    return result_str, molecule(pdb_path, residue_scores, segment), prediction_file
 
-def molecule(input_pdb, scores=None, segment='A'):
+def molecule(input_pdb, residue_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:
+    if residue_scores is not None:
+        # Sort residues based on their scores
+        high_score_residues = [resi for resi, score in residue_scores if score > 0.75]
+        mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75]
+        
         high_score_script = """
         // Reset all styles first
         viewer.getModel(0).setStyle({}, {});
@@ -127,16 +138,16 @@ def molecule(input_pdb, scores=None, segment='A'):
             {"stick": {"color": "red"}}
         );
 
-        // Highlight high-scoring residues only for the selected chain
-        let highScoreResidues2 = [%s];
+        // Highlight medium-scoring residues only for the selected chain
+        let midScoreResidues = [%s];
         viewer.getModel(0).setStyle(
-            {"chain": "%s", "resi": highScoreResidues2}, 
+            {"chain": "%s", "resi": midScoreResidues}, 
             {"stick": {"color": "orange"}}
         );
         """ % (segment, 
-               ", ".join(str(i+1) for i, score in enumerate(scores) if score > 0.8),
+               ", ".join(str(resi) for resi in high_score_residues),
                segment,
-              ", ".join(str(i+1) for i, score in enumerate(scores) if (score > 0.5) and (score < 0.8)),
+               ", ".join(str(resi) for resi in mid_score_residues),
                segment)
     
     html_content = f"""
@@ -179,7 +190,7 @@ def molecule(input_pdb, scores=None, segment='A'):
                     function(atom, viewer, event, container) {{
                         if (!atom.label) {{
                             atom.label = viewer.addLabel(
-                                atom.resn + ":" + atom.atom, 
+                                atom.resn + ":" +atom.resi + ":" + atom.atom, 
                                 {{
                                     position: atom, 
                                     backgroundColor: 'mintcream', 
@@ -246,8 +257,8 @@ with gr.Blocks() as demo:
     gr.Markdown("## Examples")
     gr.Examples(
         examples=[
-            ["2IWI", "A"],
-            ["7RPZ", "B"],
+            ["7RPZ", "A"],
+            ["2IWI", "B"],
             ["3TJN", "C"]
         ],
         inputs=[pdb_input, segment_input],

.ipynb_checkpoints/test2-checkpoint.ipynb

@@ -473,7 +473,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 1,
    "id": "d62be1b5-762e-4b69-aed4-e4ba2a44482f",
    "metadata": {},
    "outputs": [
@@ -481,7 +481,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "* Running on local URL:  http://127.0.0.1:7867\n",
+      "* Running on local URL:  http://127.0.0.1:7860\n",
       "\n",
       "To create a public link, set `share=True` in `launch()`.\n"
      ]
@@ -489,7 +489,7 @@
     {
      "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>"
+       "<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>"
@@ -502,7 +502,7 @@
      "data": {
       "text/plain": []
      },
-     "execution_count": 11,
+     "execution_count": 1,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -647,7 +647,7 @@
     "                    function(atom, viewer, event, container) {{\n",
     "                        if (!atom.label) {{\n",
     "                            atom.label = viewer.addLabel(\n",
-    "                                atom.resn + \":\" + atom.atom, \n",
+    "                                atom.resn + \":\" +atom.resi + \":\" + atom.atom, \n",
     "                                {{\n",
     "                                    position: atom, \n",
     "                                    backgroundColor: 'mintcream', \n",
@@ -727,16 +727,294 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "id": "30f35243-852f-4771-9a4b-5cdd198552b5",
    "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"
+    },
+    {
+     "data": {
+      "text/plain": []
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import gradio as gr\n",
+    "import requests\n",
+    "from Bio.PDB import PDBParser\n",
+    "import numpy as np\n",
+    "import os\n",
+    "from gradio_molecule3d import Molecule3D\n",
+    "\n",
+    "def read_mol(pdb_path):\n",
+    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
+    "    with open(pdb_path, 'r') as f:\n",
+    "        return f.read()\n",
+    "\n",
+    "def fetch_pdb(pdb_id):\n",
+    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
+    "    pdb_path = f'{pdb_id}.pdb'\n",
+    "    response = requests.get(pdb_url)\n",
+    "    if response.status_code == 200:\n",
+    "        with open(pdb_path, 'wb') as f:\n",
+    "            f.write(response.content)\n",
+    "        return pdb_path\n",
+    "    else:\n",
+    "        return None\n",
+    "\n",
+    "def process_pdb(pdb_id, segment):\n",
+    "    pdb_path = fetch_pdb(pdb_id)\n",
+    "    if not pdb_path:\n",
+    "        return \"Failed to fetch PDB file\", None, None\n",
+    "    \n",
+    "    parser = PDBParser(QUIET=1)\n",
+    "    structure = parser.get_structure('protein', pdb_path)\n",
+    "    \n",
+    "    try:\n",
+    "        chain = structure[0][segment]\n",
+    "    except KeyError:\n",
+    "        return \"Invalid Chain ID\", None, None\n",
+    "    \n",
+    "    # Comprehensive amino acid mapping\n",
+    "    aa_dict = {\n",
+    "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
+    "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
+    "        'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',\n",
+    "        'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',\n",
+    "        'MSE': 'M', 'SEP': 'S', 'TPO': 'T', 'CSO': 'C', 'PTR': 'Y', 'HYP': 'P'\n",
+    "    }\n",
+    "    \n",
+    "    # Exclude non-amino acid residues and create a list of (resi, score) pairs\n",
+    "    sequence = [\n",
+    "        (res.id[1], res) for res in chain\n",
+    "        if res.get_resname().strip() in aa_dict\n",
+    "    ]\n",
+    "\n",
+    "    random_scores = np.random.rand(len(sequence))\n",
+    "    \n",
+    "    # Zip residues with scores to track the residue ID and score\n",
+    "    residue_scores = [(resi, score) for (resi, _), score in zip(sequence, random_scores)]\n",
+    "    \n",
+    "    result_str = \"\\n\".join(\n",
+    "        f\"{aa_dict[chain[resi].get_resname()]} {resi} {score:.2f}\"\n",
+    "        for resi, score in residue_scores\n",
+    "    )\n",
+    "    \n",
+    "    # Save the predictions to a file\n",
+    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
+    "    with open(prediction_file, \"w\") as f:\n",
+    "        f.write(result_str)\n",
+    "    \n",
+    "    return result_str, molecule(pdb_path, residue_scores, segment), prediction_file\n",
+    "\n",
+    "def molecule(input_pdb, residue_scores=None, segment='A'):\n",
+    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
+    "    \n",
+    "    # Prepare high-scoring residues script if scores are provided\n",
+    "    high_score_script = \"\"\n",
+    "    if residue_scores is not None:\n",
+    "        # Sort residues based on their scores\n",
+    "        high_score_residues = [resi for resi, score in residue_scores if score > 0.9]\n",
+    "        mid_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 0.9]\n",
+    "        \n",
+    "        high_score_script = \"\"\"\n",
+    "        // Reset all styles first\n",
+    "        viewer.getModel(0).setStyle({}, {});\n",
+    "        \n",
+    "        // Show only the selected chain\n",
+    "        viewer.getModel(0).setStyle(\n",
+    "            {\"chain\": \"%s\"}, \n",
+    "            { cartoon: {colorscheme:\"whiteCarbon\"} }\n",
+    "        );\n",
+    "        \n",
+    "        // Highlight high-scoring residues only for the selected chain\n",
+    "        let highScoreResidues = [%s];\n",
+    "        viewer.getModel(0).setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": highScoreResidues}, \n",
+    "            {\"stick\": {\"color\": \"red\"}}\n",
+    "        );\n",
+    "\n",
+    "        // Highlight medium-scoring residues only for the selected chain\n",
+    "        let midScoreResidues = [%s];\n",
+    "        viewer.getModel(0).setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": midScoreResidues}, \n",
+    "            {\"stick\": {\"color\": \"orange\"}}\n",
+    "        );\n",
+    "        \"\"\" % (segment, \n",
+    "               \", \".join(str(resi) for resi in high_score_residues),\n",
+    "               segment,\n",
+    "               \", \".join(str(resi) for resi in mid_score_residues),\n",
+    "               segment)\n",
+    "    \n",
+    "    html_content = f\"\"\"\n",
+    "    <!DOCTYPE html>\n",
+    "    <html>\n",
+    "    <head>    \n",
+    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
+    "        <style>\n",
+    "        .mol-container {{\n",
+    "            width: 100%;\n",
+    "            height: 700px;\n",
+    "            position: relative;\n",
+    "        }}\n",
+    "        </style>\n",
+    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
+    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
+    "    </head>\n",
+    "    <body>\n",
+    "        <div id=\"container\" class=\"mol-container\"></div>\n",
+    "        <script>\n",
+    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
+    "            $(document).ready(function () {{\n",
+    "                let element = $(\"#container\");\n",
+    "                let config = {{ backgroundColor: \"white\" }};\n",
+    "                let viewer = $3Dmol.createViewer(element, config);\n",
+    "                viewer.addModel(pdb, \"pdb\");\n",
+    "                \n",
+    "                // Reset all styles and show only selected chain\n",
+    "                viewer.getModel(0).setStyle(\n",
+    "                    {{\"chain\": \"{segment}\"}}, \n",
+    "                    {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }}\n",
+    "                );\n",
+    "                \n",
+    "                {high_score_script}\n",
+    "                \n",
+    "                // Add hover functionality\n",
+    "                viewer.setHoverable(\n",
+    "                    {{}}, \n",
+    "                    true, \n",
+    "                    function(atom, viewer, event, container) {{\n",
+    "                        if (!atom.label) {{\n",
+    "                            atom.label = viewer.addLabel(\n",
+    "                                atom.resn + \":\" +atom.resi + \":\" + atom.atom, \n",
+    "                                {{\n",
+    "                                    position: atom, \n",
+    "                                    backgroundColor: 'mintcream', \n",
+    "                                    fontColor: 'black',\n",
+    "                                    fontSize: 12,\n",
+    "                                    padding: 2\n",
+    "                                }}\n",
+    "                            );\n",
+    "                        }}\n",
+    "                    }},\n",
+    "                    function(atom, viewer) {{\n",
+    "                        if (atom.label) {{\n",
+    "                            viewer.removeLabel(atom.label);\n",
+    "                            delete atom.label;\n",
+    "                        }}\n",
+    "                    }}\n",
+    "                );\n",
+    "                \n",
+    "                viewer.zoomTo();\n",
+    "                viewer.render();\n",
+    "                viewer.zoom(0.8, 2000);\n",
+    "            }});\n",
+    "        </script>\n",
+    "    </body>\n",
+    "    </html>\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    # Return the HTML content within an iframe safely encoded for special characters\n",
+    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
+    "\n",
+    "reps =    [\n",
+    "        {\n",
+    "          \"model\": 0,\n",
+    "          \"style\": \"cartoon\",\n",
+    "          \"color\": \"whiteCarbon\",\n",
+    "          \"residue_range\": \"\",\n",
+    "          \"around\": 0,\n",
+    "          \"byres\": False,\n",
+    "        }\n",
+    "    ]\n",
+    "\n",
+    "# Gradio UI\n",
+    "with gr.Blocks() as demo:\n",
+    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
+    "    with gr.Row():\n",
+    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
+    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
+    "\n",
+    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
+    "\n",
+    "    with gr.Row():\n",
+    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
+    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
+    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
+    "\n",
+    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
+    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
+    "    download_output = gr.File(label=\"Download Predictions\")\n",
+    "    \n",
+    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)\n",
+    "    \n",
+    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
+    "    \n",
+    "    gr.Markdown(\"## Examples\")\n",
+    "    gr.Examples(\n",
+    "        examples=[\n",
+    "            [\"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": "6f17feec-0347-4f9d-acd4-ae681c3ed425",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "63201f38-adde-4b12-a8d3-f23474d045cf",
+   "metadata": {},
    "outputs": [],
    "source": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5eca6754-4aa1-463f-881a-25d2a0d6bb5b",
+   "id": "5ccbf398-5ef2-4955-98db-99f904f8daa4",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4c61bac4-4f2e-4f4a-aa1f-30dca209747c",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -809,7 +1087,7 @@
     "    except KeyError:\n",
     "        return \"Invalid Chain ID\", None, None\n",
     "    \n",
-    "    # Comprehensive amino acid mapping\n",
+    "    \n",
     "    aa_dict = {\n",
     "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
     "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
@@ -819,9 +1097,14 @@
     "    }\n",
     "    \n",
     "    # Exclude non-amino acid residues\n",
-    "    sequence = [\n",
-    "        residue for residue in chain \n",
+    "    sequence = \"\".join(\n",
+    "        aa_dict[residue.get_resname().strip()] \n",
+    "        for residue in chain \n",
     "        if residue.get_resname().strip() in aa_dict\n",
+    "    )\n",
+    "    sequence2 = [\n",
+    "        (res.id[1], res) for res in chain\n",
+    "        if res.get_resname().strip() in aa_dict\n",
     "    ]\n",
     "    \n",
     "    # Prepare input for model prediction\n",
@@ -833,24 +1116,31 @@
     "    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",
+    "    # Zip residues with scores to track the residue ID and score\n",
+    "    residue_scores = [(resi, score) for (resi, _), score in zip(sequence2, normalized_scores)]\n",
+    "    \n",
+    "    result_str = \"\\n\".join([\n",
+    "        f\"{res.get_resname()} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}\" \n",
+    "        for i, res in enumerate(chain) if res.get_resname().strip() in aa_dict\n",
+    "    ])\n",
     "    \n",
     "    # Save the predictions to a file\n",
     "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
     "    with open(prediction_file, \"w\") as f:\n",
     "        f.write(result_str)\n",
     "    \n",
-    "    return result_str, molecule(pdb_path, random_scores, segment), prediction_file\n",
+    "    return result_str, molecule(pdb_path, residue_scores, segment), prediction_file\n",
     "\n",
-    "def molecule(input_pdb, scores=None, segment='A'):\n",
+    "def molecule(input_pdb, residue_scores=None, segment='A'):\n",
     "    mol = read_mol(input_pdb)  # Read PDB file content\n",
     "    \n",
     "    # Prepare high-scoring residues script if scores are provided\n",
     "    high_score_script = \"\"\n",
-    "    if scores is not None:\n",
+    "    if residue_scores is not None:\n",
+    "        # Sort residues based on their scores\n",
+    "        high_score_residues = [resi for resi, score in residue_scores if score > 0.75]\n",
+    "        mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75]\n",
+    "        \n",
     "        high_score_script = \"\"\"\n",
     "        // Reset all styles first\n",
     "        viewer.getModel(0).setStyle({}, {});\n",
@@ -868,16 +1158,16 @@
     "            {\"stick\": {\"color\": \"red\"}}\n",
     "        );\n",
     "\n",
-    "        // Highlight high-scoring residues only for the selected chain\n",
-    "        let highScoreResidues2 = [%s];\n",
+    "        // Highlight medium-scoring residues only for the selected chain\n",
+    "        let midScoreResidues = [%s];\n",
     "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": highScoreResidues2}, \n",
+    "            {\"chain\": \"%s\", \"resi\": midScoreResidues}, \n",
     "            {\"stick\": {\"color\": \"orange\"}}\n",
     "        );\n",
     "        \"\"\" % (segment, \n",
-    "               \", \".join(str(i+1) for i, score in enumerate(scores) if score > 0.8),\n",
+    "               \", \".join(str(resi) for resi in high_score_residues),\n",
     "               segment,\n",
-    "              \", \".join(str(i+1) for i, score in enumerate(scores) if (score > 0.5) and (score < 0.8)),\n",
+    "               \", \".join(str(resi) for resi in mid_score_residues),\n",
     "               segment)\n",
     "    \n",
     "    html_content = f\"\"\"\n",
@@ -920,7 +1210,7 @@
     "                    function(atom, viewer, event, container) {{\n",
     "                        if (!atom.label) {{\n",
     "                            atom.label = viewer.addLabel(\n",
-    "                                atom.resn + \":\" + atom.atom, \n",
+    "                                atom.resn + \":\" +atom.resi + \":\" + atom.atom, \n",
     "                                {{\n",
     "                                    position: atom, \n",
     "                                    backgroundColor: 'mintcream', \n",
@@ -987,21 +1277,21 @@
     "    gr.Markdown(\"## Examples\")\n",
     "    gr.Examples(\n",
     "        examples=[\n",
-    "            [\"2IWI\", \"A\"],\n",
-    "            [\"7RPZ\", \"B\"],\n",
+    "            [\"7RPZ\", \"A\"],\n",
+    "            [\"2IWI\", \"B\"],\n",
     "            [\"3TJN\", \"C\"]\n",
     "        ],\n",
     "        inputs=[pdb_input, segment_input],\n",
     "        outputs=[predictions_output, molecule_output, download_output]\n",
     "    )\n",
     "\n",
-    "demo.launch()"
+    "demo.launch(share=True)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "95046d1c-ec7c-4e3e-8a98-1802cb09a25b",
+   "id": "b61d06ec-a4ee-4f65-925f-d2688730416a",
    "metadata": {},
    "outputs": [],
    "source": []
@@ -1009,11 +1299,18 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a37cbe6f-d57f-41e5-8ae1-38258da39d47",
+   "id": "4d67d69f-1f53-4bcc-8905-8d29384c4e20",
    "metadata": {},
    "outputs": [],
    "source": [
     "import gradio as gr\n",
+    "import requests\n",
+    "from Bio.PDB import PDBParser\n",
+    "import numpy as np\n",
+    "import os\n",
+    "from gradio_molecule3d import Molecule3D\n",
+    "\n",
+    "\n",
     "from model_loader import load_model\n",
     "\n",
     "import torch\n",
@@ -1022,8 +1319,6 @@
     "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",
@@ -1035,18 +1330,6 @@
     "\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",
@@ -1055,23 +1338,26 @@
     "model.to(device)\n",
     "model.eval()\n",
     "\n",
-    "# Function to fetch a PDB file\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_files/{pdb_id}.pdb'\n",
-    "    os.makedirs('pdb_files', exist_ok=True)\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",
-    "    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",
+    "    else:\n",
+    "        return None\n",
     "\n",
     "def process_pdb(pdb_id, segment):\n",
     "    pdb_path = fetch_pdb(pdb_id)\n",
@@ -1080,9 +1366,13 @@
     "    \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",
+    "    try:\n",
+    "        chain = structure[0][segment]\n",
+    "    except KeyError:\n",
+    "        return \"Invalid Chain ID\", None, None\n",
+    "    \n",
+    "    \n",
     "    aa_dict = {\n",
     "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
     "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
@@ -1106,67 +1396,171 @@
     "    # 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",
+    "\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",
+    "    # 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, f\"{pdb_id}_predictions.txt\"\n",
+    "    return result_str, molecule(pdb_path, normalized_scores, segment), prediction_file\n",
     "\n",
-    "reps = [{\"model\": 0, \"style\": \"cartoon\", \"color\": \"spectrum\"}]\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\")\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\",\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",
+    "        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_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",
+    "    \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\"],\n",
-    "            [\"7RPZ\"],\n",
-    "            [\"3TJN\"]\n",
+    "            [\"2IWI\", \"A\"],\n",
+    "            [\"7RPZ\", \"B\"],\n",
+    "            [\"3TJN\", \"C\"]\n",
     "        ],\n",
-    "        inputs=[pdb_input, segment_input], \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": {

2IWI_predictions.txt

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

app.py

@@ -82,6 +82,10 @@ def process_pdb(pdb_id, segment):
         for residue in chain 
         if residue.get_resname().strip() in aa_dict
     )
+    sequence2 = [
+        (res.id[1], res) for res in chain
+        if res.get_resname().strip() in aa_dict
+    ]
     
     # Prepare input for model prediction
     input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
@@ -92,6 +96,9 @@ def process_pdb(pdb_id, segment):
     scores = expit(outputs[:, 1] - outputs[:, 0])
     normalized_scores = normalize_scores(scores)
 
+    # Zip residues with scores to track the residue ID and score
+    residue_scores = [(resi, score) for (resi, _), score in zip(sequence2, normalized_scores)]
+    
     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
@@ -102,14 +109,18 @@ def process_pdb(pdb_id, segment):
     with open(prediction_file, "w") as f:
         f.write(result_str)
     
-    return result_str, molecule(pdb_path, normalized_scores, segment), prediction_file
+    return result_str, molecule(pdb_path, residue_scores, segment), prediction_file
 
-def molecule(input_pdb, scores=None, segment='A'):
+def molecule(input_pdb, residue_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:
+    if residue_scores is not None:
+        # Sort residues based on their scores
+        high_score_residues = [resi for resi, score in residue_scores if score > 0.75]
+        mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75]
+        
         high_score_script = """
         // Reset all styles first
         viewer.getModel(0).setStyle({}, {});
@@ -127,16 +138,16 @@ def molecule(input_pdb, scores=None, segment='A'):
             {"stick": {"color": "red"}}
         );
 
-        // Highlight high-scoring residues only for the selected chain
-        let highScoreResidues2 = [%s];
+        // Highlight medium-scoring residues only for the selected chain
+        let midScoreResidues = [%s];
         viewer.getModel(0).setStyle(
-            {"chain": "%s", "resi": highScoreResidues2}, 
+            {"chain": "%s", "resi": midScoreResidues}, 
             {"stick": {"color": "orange"}}
         );
         """ % (segment, 
-               ", ".join(str(i+1) for i, score in enumerate(scores) if score > 0.8),
+               ", ".join(str(resi) for resi in high_score_residues),
                segment,
-              ", ".join(str(i+1) for i, score in enumerate(scores) if (score > 0.5) and (score < 0.8)),
+               ", ".join(str(resi) for resi in mid_score_residues),
                segment)
     
     html_content = f"""
@@ -179,7 +190,7 @@ def molecule(input_pdb, scores=None, segment='A'):
                     function(atom, viewer, event, container) {{
                         if (!atom.label) {{
                             atom.label = viewer.addLabel(
-                                atom.resn + ":" + atom.atom, 
+                                atom.resn + ":" +atom.resi + ":" + atom.atom, 
                                 {{
                                     position: atom, 
                                     backgroundColor: 'mintcream', 
@@ -246,8 +257,8 @@ with gr.Blocks() as demo:
     gr.Markdown("## Examples")
     gr.Examples(
         examples=[
-            ["2IWI", "A"],
-            ["7RPZ", "B"],
+            ["7RPZ", "A"],
+            ["2IWI", "B"],
             ["3TJN", "C"]
         ],
         inputs=[pdb_input, segment_input],

test2.ipynb

@@ -473,7 +473,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 1,
    "id": "d62be1b5-762e-4b69-aed4-e4ba2a44482f",
    "metadata": {},
    "outputs": [
@@ -481,7 +481,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "* Running on local URL:  http://127.0.0.1:7867\n",
+      "* Running on local URL:  http://127.0.0.1:7860\n",
       "\n",
       "To create a public link, set `share=True` in `launch()`.\n"
      ]
@@ -489,7 +489,7 @@
     {
      "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>"
+       "<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>"
@@ -502,7 +502,7 @@
      "data": {
       "text/plain": []
      },
-     "execution_count": 11,
+     "execution_count": 1,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -647,7 +647,7 @@
     "                    function(atom, viewer, event, container) {{\n",
     "                        if (!atom.label) {{\n",
     "                            atom.label = viewer.addLabel(\n",
-    "                                atom.resn + \":\" + atom.atom, \n",
+    "                                atom.resn + \":\" +atom.resi + \":\" + atom.atom, \n",
     "                                {{\n",
     "                                    position: atom, \n",
     "                                    backgroundColor: 'mintcream', \n",
@@ -727,16 +727,294 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "id": "30f35243-852f-4771-9a4b-5cdd198552b5",
    "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"
+    },
+    {
+     "data": {
+      "text/plain": []
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import gradio as gr\n",
+    "import requests\n",
+    "from Bio.PDB import PDBParser\n",
+    "import numpy as np\n",
+    "import os\n",
+    "from gradio_molecule3d import Molecule3D\n",
+    "\n",
+    "def read_mol(pdb_path):\n",
+    "    \"\"\"Read PDB file and return its content as a string\"\"\"\n",
+    "    with open(pdb_path, 'r') as f:\n",
+    "        return f.read()\n",
+    "\n",
+    "def fetch_pdb(pdb_id):\n",
+    "    pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'\n",
+    "    pdb_path = f'{pdb_id}.pdb'\n",
+    "    response = requests.get(pdb_url)\n",
+    "    if response.status_code == 200:\n",
+    "        with open(pdb_path, 'wb') as f:\n",
+    "            f.write(response.content)\n",
+    "        return pdb_path\n",
+    "    else:\n",
+    "        return None\n",
+    "\n",
+    "def process_pdb(pdb_id, segment):\n",
+    "    pdb_path = fetch_pdb(pdb_id)\n",
+    "    if not pdb_path:\n",
+    "        return \"Failed to fetch PDB file\", None, None\n",
+    "    \n",
+    "    parser = PDBParser(QUIET=1)\n",
+    "    structure = parser.get_structure('protein', pdb_path)\n",
+    "    \n",
+    "    try:\n",
+    "        chain = structure[0][segment]\n",
+    "    except KeyError:\n",
+    "        return \"Invalid Chain ID\", None, None\n",
+    "    \n",
+    "    # Comprehensive amino acid mapping\n",
+    "    aa_dict = {\n",
+    "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
+    "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
+    "        'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',\n",
+    "        'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y',\n",
+    "        'MSE': 'M', 'SEP': 'S', 'TPO': 'T', 'CSO': 'C', 'PTR': 'Y', 'HYP': 'P'\n",
+    "    }\n",
+    "    \n",
+    "    # Exclude non-amino acid residues and create a list of (resi, score) pairs\n",
+    "    sequence = [\n",
+    "        (res.id[1], res) for res in chain\n",
+    "        if res.get_resname().strip() in aa_dict\n",
+    "    ]\n",
+    "\n",
+    "    random_scores = np.random.rand(len(sequence))\n",
+    "    \n",
+    "    # Zip residues with scores to track the residue ID and score\n",
+    "    residue_scores = [(resi, score) for (resi, _), score in zip(sequence, random_scores)]\n",
+    "    \n",
+    "    result_str = \"\\n\".join(\n",
+    "        f\"{aa_dict[chain[resi].get_resname()]} {resi} {score:.2f}\"\n",
+    "        for resi, score in residue_scores\n",
+    "    )\n",
+    "    \n",
+    "    # Save the predictions to a file\n",
+    "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
+    "    with open(prediction_file, \"w\") as f:\n",
+    "        f.write(result_str)\n",
+    "    \n",
+    "    return result_str, molecule(pdb_path, residue_scores, segment), prediction_file\n",
+    "\n",
+    "def molecule(input_pdb, residue_scores=None, segment='A'):\n",
+    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
+    "    \n",
+    "    # Prepare high-scoring residues script if scores are provided\n",
+    "    high_score_script = \"\"\n",
+    "    if residue_scores is not None:\n",
+    "        # Sort residues based on their scores\n",
+    "        high_score_residues = [resi for resi, score in residue_scores if score > 0.9]\n",
+    "        mid_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 0.9]\n",
+    "        \n",
+    "        high_score_script = \"\"\"\n",
+    "        // Reset all styles first\n",
+    "        viewer.getModel(0).setStyle({}, {});\n",
+    "        \n",
+    "        // Show only the selected chain\n",
+    "        viewer.getModel(0).setStyle(\n",
+    "            {\"chain\": \"%s\"}, \n",
+    "            { cartoon: {colorscheme:\"whiteCarbon\"} }\n",
+    "        );\n",
+    "        \n",
+    "        // Highlight high-scoring residues only for the selected chain\n",
+    "        let highScoreResidues = [%s];\n",
+    "        viewer.getModel(0).setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": highScoreResidues}, \n",
+    "            {\"stick\": {\"color\": \"red\"}}\n",
+    "        );\n",
+    "\n",
+    "        // Highlight medium-scoring residues only for the selected chain\n",
+    "        let midScoreResidues = [%s];\n",
+    "        viewer.getModel(0).setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": midScoreResidues}, \n",
+    "            {\"stick\": {\"color\": \"orange\"}}\n",
+    "        );\n",
+    "        \"\"\" % (segment, \n",
+    "               \", \".join(str(resi) for resi in high_score_residues),\n",
+    "               segment,\n",
+    "               \", \".join(str(resi) for resi in mid_score_residues),\n",
+    "               segment)\n",
+    "    \n",
+    "    html_content = f\"\"\"\n",
+    "    <!DOCTYPE html>\n",
+    "    <html>\n",
+    "    <head>    \n",
+    "        <meta http-equiv=\"content-type\" content=\"text/html; charset=UTF-8\" />\n",
+    "        <style>\n",
+    "        .mol-container {{\n",
+    "            width: 100%;\n",
+    "            height: 700px;\n",
+    "            position: relative;\n",
+    "        }}\n",
+    "        </style>\n",
+    "        <script src=\"https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js\"></script>\n",
+    "        <script src=\"https://3Dmol.csb.pitt.edu/build/3Dmol-min.js\"></script>\n",
+    "    </head>\n",
+    "    <body>\n",
+    "        <div id=\"container\" class=\"mol-container\"></div>\n",
+    "        <script>\n",
+    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
+    "            $(document).ready(function () {{\n",
+    "                let element = $(\"#container\");\n",
+    "                let config = {{ backgroundColor: \"white\" }};\n",
+    "                let viewer = $3Dmol.createViewer(element, config);\n",
+    "                viewer.addModel(pdb, \"pdb\");\n",
+    "                \n",
+    "                // Reset all styles and show only selected chain\n",
+    "                viewer.getModel(0).setStyle(\n",
+    "                    {{\"chain\": \"{segment}\"}}, \n",
+    "                    {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }}\n",
+    "                );\n",
+    "                \n",
+    "                {high_score_script}\n",
+    "                \n",
+    "                // Add hover functionality\n",
+    "                viewer.setHoverable(\n",
+    "                    {{}}, \n",
+    "                    true, \n",
+    "                    function(atom, viewer, event, container) {{\n",
+    "                        if (!atom.label) {{\n",
+    "                            atom.label = viewer.addLabel(\n",
+    "                                atom.resn + \":\" +atom.resi + \":\" + atom.atom, \n",
+    "                                {{\n",
+    "                                    position: atom, \n",
+    "                                    backgroundColor: 'mintcream', \n",
+    "                                    fontColor: 'black',\n",
+    "                                    fontSize: 12,\n",
+    "                                    padding: 2\n",
+    "                                }}\n",
+    "                            );\n",
+    "                        }}\n",
+    "                    }},\n",
+    "                    function(atom, viewer) {{\n",
+    "                        if (atom.label) {{\n",
+    "                            viewer.removeLabel(atom.label);\n",
+    "                            delete atom.label;\n",
+    "                        }}\n",
+    "                    }}\n",
+    "                );\n",
+    "                \n",
+    "                viewer.zoomTo();\n",
+    "                viewer.render();\n",
+    "                viewer.zoom(0.8, 2000);\n",
+    "            }});\n",
+    "        </script>\n",
+    "    </body>\n",
+    "    </html>\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    # Return the HTML content within an iframe safely encoded for special characters\n",
+    "    return f'<iframe width=\"100%\" height=\"700\" srcdoc=\"{html_content.replace(chr(34), \"&quot;\").replace(chr(39), \"&#39;\")}\"></iframe>'\n",
+    "\n",
+    "reps =    [\n",
+    "        {\n",
+    "          \"model\": 0,\n",
+    "          \"style\": \"cartoon\",\n",
+    "          \"color\": \"whiteCarbon\",\n",
+    "          \"residue_range\": \"\",\n",
+    "          \"around\": 0,\n",
+    "          \"byres\": False,\n",
+    "        }\n",
+    "    ]\n",
+    "\n",
+    "# Gradio UI\n",
+    "with gr.Blocks() as demo:\n",
+    "    gr.Markdown(\"# Protein Binding Site Prediction (Random Scores)\")\n",
+    "    with gr.Row():\n",
+    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
+    "        visualize_btn = gr.Button(\"Visualize Structure\")\n",
+    "\n",
+    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
+    "\n",
+    "    with gr.Row():\n",
+    "        pdb_input = gr.Textbox(value=\"2IWI\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
+    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID\", placeholder=\"Enter Chain ID here...\")\n",
+    "        prediction_btn = gr.Button(\"Predict Random Binding Site Scores\")\n",
+    "\n",
+    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
+    "    predictions_output = gr.Textbox(label=\"Binding Site Predictions\")\n",
+    "    download_output = gr.File(label=\"Download Predictions\")\n",
+    "    \n",
+    "    visualize_btn.click(fetch_pdb, inputs=[pdb_input], outputs=molecule_output2)\n",
+    "    \n",
+    "    prediction_btn.click(process_pdb, inputs=[pdb_input, segment_input], outputs=[predictions_output, molecule_output, download_output])\n",
+    "    \n",
+    "    gr.Markdown(\"## Examples\")\n",
+    "    gr.Examples(\n",
+    "        examples=[\n",
+    "            [\"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": "6f17feec-0347-4f9d-acd4-ae681c3ed425",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "63201f38-adde-4b12-a8d3-f23474d045cf",
+   "metadata": {},
    "outputs": [],
    "source": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5eca6754-4aa1-463f-881a-25d2a0d6bb5b",
+   "id": "5ccbf398-5ef2-4955-98db-99f904f8daa4",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4c61bac4-4f2e-4f4a-aa1f-30dca209747c",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -809,7 +1087,7 @@
     "    except KeyError:\n",
     "        return \"Invalid Chain ID\", None, None\n",
     "    \n",
-    "    # Comprehensive amino acid mapping\n",
+    "    \n",
     "    aa_dict = {\n",
     "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
     "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
@@ -819,9 +1097,14 @@
     "    }\n",
     "    \n",
     "    # Exclude non-amino acid residues\n",
-    "    sequence = [\n",
-    "        residue for residue in chain \n",
+    "    sequence = \"\".join(\n",
+    "        aa_dict[residue.get_resname().strip()] \n",
+    "        for residue in chain \n",
     "        if residue.get_resname().strip() in aa_dict\n",
+    "    )\n",
+    "    sequence2 = [\n",
+    "        (res.id[1], res) for res in chain\n",
+    "        if res.get_resname().strip() in aa_dict\n",
     "    ]\n",
     "    \n",
     "    # Prepare input for model prediction\n",
@@ -833,24 +1116,31 @@
     "    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",
+    "    # Zip residues with scores to track the residue ID and score\n",
+    "    residue_scores = [(resi, score) for (resi, _), score in zip(sequence2, normalized_scores)]\n",
+    "    \n",
+    "    result_str = \"\\n\".join([\n",
+    "        f\"{res.get_resname()} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}\" \n",
+    "        for i, res in enumerate(chain) if res.get_resname().strip() in aa_dict\n",
+    "    ])\n",
     "    \n",
     "    # Save the predictions to a file\n",
     "    prediction_file = f\"{pdb_id}_predictions.txt\"\n",
     "    with open(prediction_file, \"w\") as f:\n",
     "        f.write(result_str)\n",
     "    \n",
-    "    return result_str, molecule(pdb_path, random_scores, segment), prediction_file\n",
+    "    return result_str, molecule(pdb_path, residue_scores, segment), prediction_file\n",
     "\n",
-    "def molecule(input_pdb, scores=None, segment='A'):\n",
+    "def molecule(input_pdb, residue_scores=None, segment='A'):\n",
     "    mol = read_mol(input_pdb)  # Read PDB file content\n",
     "    \n",
     "    # Prepare high-scoring residues script if scores are provided\n",
     "    high_score_script = \"\"\n",
-    "    if scores is not None:\n",
+    "    if residue_scores is not None:\n",
+    "        # Sort residues based on their scores\n",
+    "        high_score_residues = [resi for resi, score in residue_scores if score > 0.75]\n",
+    "        mid_score_residues = [resi for resi, score in residue_scores if 0.5 < score <= 0.75]\n",
+    "        \n",
     "        high_score_script = \"\"\"\n",
     "        // Reset all styles first\n",
     "        viewer.getModel(0).setStyle({}, {});\n",
@@ -868,16 +1158,16 @@
     "            {\"stick\": {\"color\": \"red\"}}\n",
     "        );\n",
     "\n",
-    "        // Highlight high-scoring residues only for the selected chain\n",
-    "        let highScoreResidues2 = [%s];\n",
+    "        // Highlight medium-scoring residues only for the selected chain\n",
+    "        let midScoreResidues = [%s];\n",
     "        viewer.getModel(0).setStyle(\n",
-    "            {\"chain\": \"%s\", \"resi\": highScoreResidues2}, \n",
+    "            {\"chain\": \"%s\", \"resi\": midScoreResidues}, \n",
     "            {\"stick\": {\"color\": \"orange\"}}\n",
     "        );\n",
     "        \"\"\" % (segment, \n",
-    "               \", \".join(str(i+1) for i, score in enumerate(scores) if score > 0.8),\n",
+    "               \", \".join(str(resi) for resi in high_score_residues),\n",
     "               segment,\n",
-    "              \", \".join(str(i+1) for i, score in enumerate(scores) if (score > 0.5) and (score < 0.8)),\n",
+    "               \", \".join(str(resi) for resi in mid_score_residues),\n",
     "               segment)\n",
     "    \n",
     "    html_content = f\"\"\"\n",
@@ -920,7 +1210,7 @@
     "                    function(atom, viewer, event, container) {{\n",
     "                        if (!atom.label) {{\n",
     "                            atom.label = viewer.addLabel(\n",
-    "                                atom.resn + \":\" + atom.atom, \n",
+    "                                atom.resn + \":\" +atom.resi + \":\" + atom.atom, \n",
     "                                {{\n",
     "                                    position: atom, \n",
     "                                    backgroundColor: 'mintcream', \n",
@@ -987,21 +1277,21 @@
     "    gr.Markdown(\"## Examples\")\n",
     "    gr.Examples(\n",
     "        examples=[\n",
-    "            [\"2IWI\", \"A\"],\n",
-    "            [\"7RPZ\", \"B\"],\n",
+    "            [\"7RPZ\", \"A\"],\n",
+    "            [\"2IWI\", \"B\"],\n",
     "            [\"3TJN\", \"C\"]\n",
     "        ],\n",
     "        inputs=[pdb_input, segment_input],\n",
     "        outputs=[predictions_output, molecule_output, download_output]\n",
     "    )\n",
     "\n",
-    "demo.launch()"
+    "demo.launch(share=True)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "95046d1c-ec7c-4e3e-8a98-1802cb09a25b",
+   "id": "b61d06ec-a4ee-4f65-925f-d2688730416a",
    "metadata": {},
    "outputs": [],
    "source": []
@@ -1009,11 +1299,18 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a37cbe6f-d57f-41e5-8ae1-38258da39d47",
+   "id": "4d67d69f-1f53-4bcc-8905-8d29384c4e20",
    "metadata": {},
    "outputs": [],
    "source": [
     "import gradio as gr\n",
+    "import requests\n",
+    "from Bio.PDB import PDBParser\n",
+    "import numpy as np\n",
+    "import os\n",
+    "from gradio_molecule3d import Molecule3D\n",
+    "\n",
+    "\n",
     "from model_loader import load_model\n",
     "\n",
     "import torch\n",
@@ -1022,8 +1319,6 @@
     "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",
@@ -1035,18 +1330,6 @@
     "\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",
@@ -1055,23 +1338,26 @@
     "model.to(device)\n",
     "model.eval()\n",
     "\n",
-    "# Function to fetch a PDB file\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_files/{pdb_id}.pdb'\n",
-    "    os.makedirs('pdb_files', exist_ok=True)\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",
-    "    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",
+    "    else:\n",
+    "        return None\n",
     "\n",
     "def process_pdb(pdb_id, segment):\n",
     "    pdb_path = fetch_pdb(pdb_id)\n",
@@ -1080,9 +1366,13 @@
     "    \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",
+    "    try:\n",
+    "        chain = structure[0][segment]\n",
+    "    except KeyError:\n",
+    "        return \"Invalid Chain ID\", None, None\n",
+    "    \n",
+    "    \n",
     "    aa_dict = {\n",
     "        'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',\n",
     "        'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',\n",
@@ -1106,67 +1396,171 @@
     "    # 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",
+    "\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",
+    "    # 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, f\"{pdb_id}_predictions.txt\"\n",
+    "    return result_str, molecule(pdb_path, normalized_scores, segment), prediction_file\n",
     "\n",
-    "reps = [{\"model\": 0, \"style\": \"cartoon\", \"color\": \"spectrum\"}]\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\")\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\",\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",
+    "        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_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",
+    "    \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\"],\n",
-    "            [\"7RPZ\"],\n",
-    "            [\"3TJN\"]\n",
+    "            [\"2IWI\", \"A\"],\n",
+    "            [\"7RPZ\", \"B\"],\n",
+    "            [\"3TJN\", \"C\"]\n",
     "        ],\n",
-    "        inputs=[pdb_input, segment_input], \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": {