Update

.gradio/cached_examples/148/log.csv

@@ -0,0 +1,4 @@
+PyMOL Visualization Commands,Interactive 3D Structure,Download Results,component 3,timestamp
+Failed to create chain-specific PDB: invalid literal for int() with base 10: 'THR',,,"{'visible': False, '__type__': 'update'}",2024-12-27 14:49:10.613413
+Failed to create chain-specific PDB: invalid literal for int() with base 10: 'YR',,,"{'visible': False, '__type__': 'update'}",2024-12-27 14:49:10.680800
+Failed to create chain-specific PDB: invalid literal for int() with base 10: 'LU',,,"{'visible': False, '__type__': 'update'}",2024-12-27 14:49:10.744581

.gradio/cached_examples/72/log.csv

@@ -0,0 +1,4 @@
+PyMOL Visualization Commands,Interactive 3D Structure,Download Results,component 3,timestamp
+Failed to create chain-specific PDB: invalid literal for int() with base 10: 'THR',,,"{'visible': False, '__type__': 'update'}",2024-12-27 14:47:04.996998
+Failed to create chain-specific PDB: invalid literal for int() with base 10: 'YR',,,"{'visible': False, '__type__': 'update'}",2024-12-27 14:47:05.557021
+Failed to create chain-specific PDB: invalid literal for int() with base 10: 'LU',,,"{'visible': False, '__type__': 'update'}",2024-12-27 14:47:06.032725

.ipynb_checkpoints/app-checkpoint.py

@@ -139,30 +139,6 @@ def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: lis
     
     return output_pdb
 
-def calculate_geometric_center(pdb_path: str, high_score_residues: list, chain_id: str):
-    """
-    Calculate the geometric center of high-scoring residues
-    """
-    parser = PDBParser(QUIET=True)
-    structure = parser.get_structure('protein', pdb_path)
-    
-    # Collect coordinates of CA atoms from high-scoring residues
-    coords = []
-    for model in structure:
-        for chain in model:
-            if chain.id == chain_id:
-                for residue in chain:
-                    if residue.id[1] in high_score_residues:
-                        if 'CA' in residue:  # Use alpha carbon as representative
-                            ca_atom = residue['CA']
-                            coords.append(ca_atom.coord)
-    
-    # Calculate geometric center
-    if coords:
-        center = np.mean(coords, axis=0)
-        return center
-    return None
-
 def process_pdb(pdb_id_or_file, segment):
     # Determine if input is a PDB ID or file path
     if pdb_id_or_file.endswith('.pdb'):
@@ -194,7 +170,11 @@ def process_pdb(pdb_id_or_file, segment):
     protein_residues = [res for res in chain if is_aa(res)]
     sequence = "".join(seq1(res.resname) for res in protein_residues)
     sequence_id = [res.id[1] for res in protein_residues]
-    
+
+    visualized_sequence = "".join(seq1(res.resname) for res in protein_residues)
+    if sequence != visualized_sequence:
+        raise ValueError("The visualized sequence does not match the prediction sequence")
+     
     input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
     with torch.no_grad():
         outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()
@@ -300,7 +280,6 @@ def molecule(input_pdb, residue_scores=None, segment='A'):
         class4_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.8]
         class5_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 1.0]
         
-        
         high_score_script = """
         // Load the original model and apply white cartoon style
         let chainModel = viewer.addModel(pdb, "pdb");
@@ -430,7 +409,19 @@ def molecule(input_pdb, residue_scores=None, segment='A'):
     return f'<iframe width="100%" height="700" srcdoc="{html_content.replace(chr(34), "&quot;").replace(chr(39), "&#39;")}"></iframe>'
 
 # Gradio UI
-with gr.Blocks() as demo:
+with gr.Blocks(css="""
+    /* Customize Gradio button colors */
+    #visualize-btn, #predict-btn {
+        background-color: #FF7300; /* Deep orange */
+        color: white;
+        border-radius: 5px;
+        padding: 10px;
+        font-weight: bold;
+    }
+    #visualize-btn:hover, #predict-btn:hover {
+        background-color: #CC5C00; /* Darkened orange on hover */
+    }
+""") as demo:
     gr.Markdown("# Protein Binding Site Prediction")
     
     # Mode selection
@@ -442,9 +433,9 @@ with gr.Blocks() as demo:
     )
 
     # Input components based on mode
-    pdb_input = gr.Textbox(value="4BDU", label="PDB ID", placeholder="Enter PDB ID here...")
+    pdb_input = gr.Textbox(value="2F6V", label="PDB ID", placeholder="Enter PDB ID here...")
     pdb_file = gr.File(label="Upload PDB/CIF File", visible=False)
-    visualize_btn = gr.Button("Visualize Structure")
+    visualize_btn = gr.Button("Visualize Structure", elem_id="visualize-btn")
 
     molecule_output2 = Molecule3D(label="Protein Structure", reps=[
         {
@@ -458,8 +449,9 @@ with gr.Blocks() as demo:
     ])
 
     with gr.Row():
-        segment_input = gr.Textbox(value="A", label="Chain ID", placeholder="Enter Chain ID here...")
-        prediction_btn = gr.Button("Predict Binding Site")
+        segment_input = gr.Textbox(value="A", label="Chain ID (protein)", placeholder="Enter Chain ID here...",
+        info="Choose in which chain to predict binding sites.")
+        prediction_btn = gr.Button("Predict Binding Site", elem_id="predict-btn")
 
     molecule_output = gr.HTML(label="Protein Structure")
     explanation_vis = gr.Markdown("""
@@ -533,7 +525,7 @@ with gr.Blocks() as demo:
         examples=[
             ["7RPZ", "A"],
             ["2IWI", "B"],
-            ["2F6V", "A"]
+            ["7LCJ", "R"]
         ],
         inputs=[pdb_input, segment_input],
         outputs=[predictions_output, molecule_output, download_output]

.ipynb_checkpoints/test-checkpoint.ipynb

@@ -2,23 +2,24 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 3,
-   "id": "1f8ea359-674c-4263-9c2a-7a8e7e464249",
+   "execution_count": 29,
+   "id": "e776d9d6-417e-46d4-8061-846c055e1f8a",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "* Running on local URL:  http://127.0.0.1:7862\n",
+      "* Running on local URL:  http://127.0.0.1:7873\n",
+      "* Running on public URL: https://120000a6aa9d78e04c.gradio.live\n",
       "\n",
-      "To create a public link, set `share=True` in `launch()`.\n"
+      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
      ]
     },
     {
      "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>"
+       "<div><iframe src=\"https://120000a6aa9d78e04c.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"
@@ -31,501 +32,344 @@
      "data": {
       "text/plain": []
      },
-     "execution_count": 3,
+     "execution_count": 29,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
+    "from datetime import datetime\n",
     "import gradio as gr\n",
     "import requests\n",
-    "from Bio.PDB import PDBParser\n",
-    "from gradio_molecule3d import Molecule3D\n",
+    "from Bio.PDB import PDBParser, MMCIFParser, PDBIO, Select\n",
+    "from Bio.PDB.Polypeptide import is_aa\n",
+    "from Bio.SeqUtils import seq1\n",
+    "from typing import Optional, Tuple\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",
+    "import os\n",
     "from gradio_molecule3d import Molecule3D\n",
     "\n",
+    "#from model_loader import load_model\n",
     "\n",
-    "example = Molecule3D().example_value()\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",
-    "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",
+    "import re\n",
+    "import pandas as pd\n",
+    "import copy\n",
     "\n",
+    "#import transformers\n",
+    "#from transformers import AutoTokenizer, DataCollatorForTokenClassification\n",
     "\n",
+    "#from datasets import Dataset\n",
     "\n",
-    "def predict(x):\n",
-    "    print(\"predict function\", x)\n",
-    "    print(x.name)\n",
-    "    return x\n",
+    "from scipy.special import expit\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",
+    "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",
-    "# 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",
+    "def fetch_structure(pdb_id: str, output_dir: str = \".\") -> Optional[str]:\n",
+    "    \"\"\"\n",
+    "    Fetch the structure file for a given PDB ID. Prioritizes CIF files.\n",
+    "    If a structure file already exists locally, it uses that.\n",
+    "    \"\"\"\n",
+    "    file_path = download_structure(pdb_id, output_dir)\n",
+    "    if file_path:\n",
+    "        return file_path\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",
+    "def download_structure(pdb_id: str, output_dir: str) -> Optional[str]:\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",
+    "    Attempt to download the structure file in CIF or PDB format.\n",
+    "    Returns the path to the downloaded file, or None if download fails.\n",
+    "    \"\"\"\n",
+    "    for ext in ['.cif', '.pdb']:\n",
+    "        file_path = os.path.join(output_dir, f\"{pdb_id}{ext}\")\n",
+    "        if os.path.exists(file_path):\n",
+    "            return file_path\n",
+    "        url = f\"https://files.rcsb.org/download/{pdb_id}{ext}\"\n",
+    "        try:\n",
+    "            response = requests.get(url, timeout=10)\n",
+    "            if response.status_code == 200:\n",
+    "                with open(file_path, 'wb') as f:\n",
+    "                    f.write(response.content)\n",
+    "                return file_path\n",
+    "        except Exception as e:\n",
+    "            print(f\"Download error for {pdb_id}{ext}: {e}\")\n",
+    "    return None\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",
+    "def convert_cif_to_pdb(cif_path: str, output_dir: str = \".\") -> str:\n",
+    "    \"\"\"\n",
+    "    Convert a CIF file to PDB format using BioPython and return the PDB file path.\n",
+    "    \"\"\"\n",
+    "    pdb_path = os.path.join(output_dir, os.path.basename(cif_path).replace('.cif', '.pdb'))\n",
+    "    parser = MMCIFParser(QUIET=True)\n",
+    "    structure = parser.get_structure('protein', cif_path)\n",
+    "    io = PDBIO()\n",
+    "    io.set_structure(structure)\n",
+    "    io.save(pdb_path)\n",
+    "    return pdb_path\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",
+    "def fetch_pdb(pdb_id):\n",
+    "    pdb_path = fetch_structure(pdb_id)\n",
+    "    if not pdb_path:\n",
+    "        return None\n",
+    "    _, ext = os.path.splitext(pdb_path)\n",
+    "    if ext == '.cif':\n",
+    "        pdb_path = convert_cif_to_pdb(pdb_path)\n",
+    "    return pdb_path\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",
+    "def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list, protein_residues: list) -> str:\n",
+    "    \"\"\"\n",
+    "    Create a PDB file with only the selected chain and residues, replacing B-factor with prediction scores\n",
+    "    \"\"\"\n",
+    "    # Read the original PDB file\n",
+    "    parser = PDBParser(QUIET=True)\n",
+    "    structure = parser.get_structure('protein', input_pdb)\n",
+    "    \n",
+    "    # Prepare a new structure with only the specified chain and selected residues\n",
+    "    output_pdb = f\"{os.path.splitext(input_pdb)[0]}_{chain_id}_predictions_scores.pdb\"\n",
+    "    \n",
+    "    # Create scores dictionary for easy lookup\n",
+    "    scores_dict = {resi: score for resi, score in residue_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",
+    "    # Create a custom Select class\n",
+    "    class ResidueSelector(Select):\n",
+    "        def __init__(self, chain_id, selected_residues, scores_dict):\n",
+    "            self.chain_id = chain_id\n",
+    "            self.selected_residues = selected_residues\n",
+    "            self.scores_dict = scores_dict\n",
+    "        \n",
+    "        def accept_chain(self, chain):\n",
+    "            return chain.id == self.chain_id\n",
+    "        \n",
+    "        def accept_residue(self, residue):\n",
+    "            return residue.id[1] in self.selected_residues\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",
+    "        def accept_atom(self, atom):\n",
+    "            if atom.parent.id[1] in self.scores_dict:\n",
+    "                atom.bfactor = np.absolute(1-self.scores_dict[atom.parent.id[1]]) * 100\n",
+    "            return True\n",
+    "\n",
+    "    # Prepare output PDB with selected chain and residues, modified B-factors\n",
+    "    io = PDBIO()\n",
+    "    selector = ResidueSelector(chain_id, [res.id[1] for res in protein_residues], scores_dict)\n",
+    "    \n",
+    "    io.set_structure(structure[0])\n",
+    "    io.save(output_pdb, selector)\n",
+    "    \n",
+    "    return output_pdb\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",
+    "def process_pdb(pdb_id_or_file, segment):\n",
+    "    # Determine if input is a PDB ID or file path\n",
+    "    if pdb_id_or_file.endswith('.pdb'):\n",
+    "        pdb_path = pdb_id_or_file\n",
+    "        pdb_id = os.path.splitext(os.path.basename(pdb_path))[0]\n",
+    "    else:\n",
+    "        pdb_id = pdb_id_or_file\n",
+    "        pdb_path = fetch_pdb(pdb_id)\n",
+    "    \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",
+    "    # Determine the file format and choose the appropriate parser\n",
+    "    _, ext = os.path.splitext(pdb_path)\n",
+    "    parser = MMCIFParser(QUIET=True) if ext == '.cif' else PDBParser(QUIET=True)\n",
     "    \n",
     "    try:\n",
+    "        # Parse the structure file\n",
+    "        structure = parser.get_structure('protein', pdb_path)\n",
+    "    except Exception as e:\n",
+    "        return f\"Error parsing structure file: {e}\", None, None\n",
+    "    \n",
+    "    # Extract the specified chain\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",
+    "    protein_residues = [res for res in chain if is_aa(res)]\n",
+    "    sequence = \"\".join(seq1(res.resname) for res in protein_residues)\n",
+    "    sequence_id = [res.id[1] for res in protein_residues]\n",
+    "\n",
+    "    visualized_sequence = \"\".join(seq1(res.resname) for res in protein_residues)\n",
+    "    if sequence != visualized_sequence:\n",
+    "        raise ValueError(\"The visualized sequence does not match the prediction sequence\")\n",
+    "        \n",
+    "    scores = np.random.rand(len(sequence))\n",
+    "    normalized_scores = normalize_scores(scores)\n",
     "    \n",
-    "    return result_str, molecule(pdb_path), prediction_file\n",
+    "    # Zip residues with scores to track the residue ID and score\n",
+    "    residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]\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",
+    "    # Define the score brackets\n",
+    "    score_brackets = {\n",
+    "        \"0.0-0.2\": (0.0, 0.2),\n",
+    "        \"0.2-0.4\": (0.2, 0.4),\n",
+    "        \"0.4-0.6\": (0.4, 0.6),\n",
+    "        \"0.6-0.8\": (0.6, 0.8),\n",
+    "        \"0.8-1.0\": (0.8, 1.0)\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",
+    "    # Initialize a dictionary to store residues by bracket\n",
+    "    residues_by_bracket = {bracket: [] for bracket in score_brackets}\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",
+    "    # Categorize residues into brackets\n",
+    "    for resi, score in residue_scores:\n",
+    "        for bracket, (lower, upper) in score_brackets.items():\n",
+    "            if lower <= score < upper:\n",
+    "                residues_by_bracket[bracket].append(resi)\n",
+    "                break\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",
+    "    # Preparing the result string\n",
+    "    current_time = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+    "    result_str = f\"Prediction for PDB: {pdb_id}, Chain: {segment}\\nDate: {current_time}\\n\\n\"\n",
+    "    result_str += \"Residues by Score Brackets:\\n\\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",
+    "    # Add residues for each bracket\n",
+    "    for bracket, residues in residues_by_bracket.items():\n",
+    "        result_str += f\"Bracket {bracket}:\\n\"\n",
+    "        result_str += \"Columns: Residue Name, Residue Number, One-letter Code, Normalized Score\\n\"\n",
+    "        result_str += \"\\n\".join([\n",
+    "            f\"{res.resname} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}\" \n",
+    "            for i, res in enumerate(protein_residues) if res.id[1] in residues\n",
+    "        ])\n",
+    "        result_str += \"\\n\\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",
+    "    # Create chain-specific PDB with scores in B-factor\n",
+    "    scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores, protein_residues)\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",
+    "    # Molecule visualization with updated script with color mapping\n",
+    "    mol_vis = molecule(pdb_path, residue_scores, segment)#, color_map)\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",
+    "    # Improved PyMOL command suggestions\n",
+    "    current_time = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+    "    pymol_commands = f\"Prediction for PDB: {pdb_id}, Chain: {segment}\\nDate: {current_time}\\n\\n\"\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",
+    "    pymol_commands += f\"\"\"\n",
+    "    # PyMOL Visualization Commands\n",
+    "    load {os.path.abspath(pdb_path)}, protein\n",
+    "    hide everything, all\n",
+    "    show cartoon, chain {segment}\n",
+    "    color white, chain {segment}\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    # Define colors for each score bracket\n",
+    "    bracket_colors = {\n",
+    "        \"0.0-0.2\": \"white\",\n",
+    "        \"0.2-0.4\": \"lightorange\",\n",
+    "        \"0.4-0.6\": \"orange\",\n",
+    "        \"0.6-0.8\": \"orangered\",\n",
+    "        \"0.8-1.0\": \"red\"\n",
+    "    }\n",
+    "    \n",
+    "    # Add PyMOL commands for each score bracket\n",
+    "    for bracket, residues in residues_by_bracket.items():\n",
+    "        if residues:  # Only add commands if there are residues in this bracket\n",
+    "            color = bracket_colors[bracket]\n",
+    "            resi_list = '+'.join(map(str, residues))\n",
+    "            pymol_commands += f\"\"\"\n",
+    "    select bracket_{bracket.replace('.', '').replace('-', '_')}, resi {resi_list} and chain {segment}\n",
+    "    show sticks, bracket_{bracket.replace('.', '').replace('-', '_')}\n",
+    "    color {color}, bracket_{bracket.replace('.', '').replace('-', '_')}\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    # Create prediction and scored PDB files\n",
+    "    prediction_file = f\"{pdb_id}_binding_site_residues.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",
+    "    return pymol_commands, mol_vis, [prediction_file,scored_pdb]\n",
     "\n",
-    "def molecule(input_pdb):\n",
+    "def molecule(input_pdb, residue_scores=None, segment='A'):\n",
+    "    # More granular scoring for visualization\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",
+    "        # Filter residues based on their scores\n",
+    "        class1_score_residues = [resi for resi, score in residue_scores if 0.0 < score <= 0.2]\n",
+    "        class2_score_residues = [resi for resi, score in residue_scores if 0.2 < score <= 0.4]\n",
+    "        class3_score_residues = [resi for resi, score in residue_scores if 0.4 < score <= 0.6]\n",
+    "        class4_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.8]\n",
+    "        class5_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 1.0]\n",
+    "        \n",
+    "        high_score_script = \"\"\"\n",
+    "        // Load the original model and apply white cartoon style\n",
+    "        let chainModel = viewer.addModel(pdb, \"pdb\");\n",
+    "        chainModel.setStyle({}, {});\n",
+    "        chainModel.setStyle(\n",
+    "            {\"chain\": \"%s\"}, \n",
+    "            {\"cartoon\": {\"color\": \"white\"}}\n",
+    "        );\n",
+    "\n",
+    "        // Create a new model for high-scoring residues and apply red sticks style\n",
+    "        let class1Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        class1Model.setStyle({}, {});\n",
+    "        class1Model.setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
+    "            {\"stick\": {\"color\": \"0xFFFFFF\", \"opacity\": 0.5}}\n",
+    "        );\n",
+    "\n",
+    "        // Create a new model for high-scoring residues and apply red sticks style\n",
+    "        let class2Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        class2Model.setStyle({}, {});\n",
+    "        class2Model.setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
+    "            {\"stick\": {\"color\": \"0xFFD580\", \"opacity\": 0.7}}\n",
+    "        );\n",
+    "\n",
+    "        // Create a new model for high-scoring residues and apply red sticks style\n",
+    "        let class3Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        class3Model.setStyle({}, {});\n",
+    "        class3Model.setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
+    "            {\"stick\": {\"color\": \"0xFFA500\", \"opacity\": 1}}\n",
+    "        );\n",
+    "\n",
+    "        // Create a new model for high-scoring residues and apply red sticks style\n",
+    "        let class4Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        class4Model.setStyle({}, {});\n",
+    "        class4Model.setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
+    "            {\"stick\": {\"color\": \"0xFF4500\", \"opacity\": 1}}\n",
+    "        );\n",
+    "\n",
+    "        // Create a new model for high-scoring residues and apply red sticks style\n",
+    "        let class5Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        class5Model.setStyle({}, {});\n",
+    "        class5Model.setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
+    "            {\"stick\": {\"color\": \"0xFF0000\", \"alpha\": 1}}\n",
+    "        );\n",
+    "\n",
+    "        \"\"\" % (\n",
+    "            segment,\n",
+    "            segment,\n",
+    "            \", \".join(str(resi) for resi in class1_score_residues),\n",
+    "            segment,\n",
+    "            \", \".join(str(resi) for resi in class2_score_residues),\n",
+    "            segment,\n",
+    "            \", \".join(str(resi) for resi in class3_score_residues),\n",
+    "            segment,\n",
+    "            \", \".join(str(resi) for resi in class4_score_residues),\n",
+    "            segment,\n",
+    "            \", \".join(str(resi) for resi in class5_score_residues)\n",
+    "        )\n",
     "    \n",
+    "    # Generate the full HTML content\n",
     "    html_content = f\"\"\"\n",
     "    <!DOCTYPE html>\n",
     "    <html>\n",
@@ -549,15 +393,33 @@
     "                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",
+    "                {high_score_script}\n",
     "                \n",
-    "                // Highlight specific histidine residues in red stick representation\n",
-    "                viewer.getModel(0).setStyle(\n",
-    "                    {{\"resn\": \"HIS\"}}, \n",
-    "                    {{\"stick\": {{\"color\": \"red\"}}}}\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: 18,\n",
+    "                                    padding: 4\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",
@@ -573,68 +435,194 @@
     "    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",
+    "with gr.Blocks(css=\"\"\"\n",
+    "    /* Customize Gradio button colors */\n",
+    "    #visualize-btn, #predict-btn {\n",
+    "        background-color: #FF7300; /* Deep orange */\n",
+    "        color: white;\n",
+    "        border-radius: 5px;\n",
+    "        padding: 10px;\n",
+    "        font-weight: bold;\n",
+    "    }\n",
+    "    #visualize-btn:hover, #predict-btn:hover {\n",
+    "        background-color: #CC5C00; /* Darkened orange on hover */\n",
+    "    }\n",
+    "\"\"\") as demo:\n",
+    "    gr.Markdown(\"# Protein Binding Site Prediction\")\n",
     "    \n",
+    "    # Mode selection\n",
+    "    mode = gr.Radio(\n",
+    "        choices=[\"PDB ID\", \"Upload File\"],\n",
+    "        value=\"PDB ID\",\n",
+    "        label=\"Input Mode\",\n",
+    "        info=\"Choose whether to input a PDB ID or upload a PDB/CIF file.\"\n",
+    "    )\n",
+    "\n",
+    "    # Input components based on mode\n",
+    "    pdb_input = gr.Textbox(value=\"2F6V\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
+    "    pdb_file = gr.File(label=\"Upload PDB/CIF File\", visible=False)\n",
+    "    visualize_btn = gr.Button(\"Visualize Structure\", elem_id=\"visualize-btn\")\n",
+    "\n",
+    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=[\n",
+    "        {\n",
+    "            \"model\": 0,\n",
+    "            \"style\": \"cartoon\",\n",
+    "            \"color\": \"whiteCarbon\",\n",
+    "            \"residue_range\": \"\",\n",
+    "            \"around\": 0,\n",
+    "            \"byres\": False,\n",
+    "        }\n",
+    "    ])\n",
+    "\n",
+    "    with gr.Row():\n",
+    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID (protein)\", placeholder=\"Enter Chain ID here...\",\n",
+    "        info=\"Choose in which chain to predict binding sites.\")\n",
+    "        prediction_btn = gr.Button(\"Predict Binding Site\", elem_id=\"predict-btn\")\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",
+    "    explanation_vis = gr.Markdown(\"\"\"\n",
+    "    Score dependent colorcoding:\n",
+    "    - 0.0-0.2: white  \n",
+    "    - 0.2–0.4: light orange  \n",
+    "    - 0.4–0.6: orange\n",
+    "    - 0.6–0.8: orangered\n",
+    "    - 0.8–1.0: red\n",
+    "    \"\"\")\n",
+    "    predictions_output = gr.Textbox(label=\"Visualize Prediction with PyMol\")\n",
+    "    gr.Markdown(\"### Download:\\n- List of predicted binding site residues\\n- PDB with score in beta factor column\")\n",
+    "    download_output = gr.File(label=\"Download Files\", file_count=\"multiple\")\n",
     "    \n",
-    "    # Update to explicitly use molecule() function for visualization\n",
+    "    def process_interface(mode, pdb_id, pdb_file, chain_id):\n",
+    "        if mode == \"PDB ID\":\n",
+    "            return process_pdb(pdb_id, chain_id)\n",
+    "        elif mode == \"Upload File\":\n",
+    "            _, ext = os.path.splitext(pdb_file.name)\n",
+    "            file_path = os.path.join('./', f\"{_}{ext}\")\n",
+    "            if ext == '.cif':\n",
+    "                pdb_path = convert_cif_to_pdb(file_path)\n",
+    "            else:\n",
+    "                pdb_path= file_path\n",
+    "            return process_pdb(pdb_path, chain_id)\n",
+    "        else:\n",
+    "            return \"Error: Invalid mode selected\", None, None\n",
+    "\n",
+    "    def fetch_interface(mode, pdb_id, pdb_file):\n",
+    "        if mode == \"PDB ID\":\n",
+    "            return fetch_pdb(pdb_id)\n",
+    "        elif mode == \"Upload File\":\n",
+    "            _, ext = os.path.splitext(pdb_file.name)\n",
+    "            file_path = os.path.join('./', f\"{_}{ext}\")\n",
+    "            #print(ext)\n",
+    "            if ext == '.cif':\n",
+    "                pdb_path = convert_cif_to_pdb(file_path)\n",
+    "            else:\n",
+    "                pdb_path= file_path\n",
+    "            #print(pdb_path)\n",
+    "            return pdb_path\n",
+    "        else:\n",
+    "            return \"Error: Invalid mode selected\"\n",
+    "\n",
+    "    def toggle_mode(selected_mode):\n",
+    "        if selected_mode == \"PDB ID\":\n",
+    "            return gr.update(visible=True), gr.update(visible=False)\n",
+    "        else:\n",
+    "            return gr.update(visible=False), gr.update(visible=True)\n",
+    "\n",
+    "    mode.change(\n",
+    "        toggle_mode,\n",
+    "        inputs=[mode],\n",
+    "        outputs=[pdb_input, pdb_file]\n",
+    "    )\n",
+    "\n",
+    "    prediction_btn.click(\n",
+    "        process_interface, \n",
+    "        inputs=[mode, pdb_input, pdb_file, segment_input], \n",
+    "        outputs=[predictions_output, molecule_output, download_output]\n",
+    "    )\n",
+    "\n",
     "    visualize_btn.click(\n",
-    "        fn=lambda pdb_id: molecule(fetch_pdb(pdb_id)), \n",
-    "        inputs=[pdb_input], \n",
-    "        outputs=molecule_output\n",
+    "        fetch_interface, \n",
+    "        inputs=[mode, pdb_input, pdb_file], \n",
+    "        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",
+    "\n",
     "    gr.Markdown(\"## Examples\")\n",
     "    gr.Examples(\n",
     "        examples=[\n",
-    "            [\"2IWI\", \"A\"],\n",
-    "            [\"7RPZ\", \"B\"],\n",
-    "            [\"3TJN\", \"C\"]\n",
+    "            [\"7RPZ\", \"A\"],\n",
+    "            [\"2IWI\", \"B\"],\n",
+    "            [\"7LCJ\", \"R\"]\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": "15527a58-c449-4da0-8fab-3baaede15e41",
+   "id": "440c87ed-45c9-4501-b208-409cbfd7858b",
    "metadata": {},
    "outputs": [],
    "source": []
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
-   "id": "9ef3e330-cb88-4c29-b84a-2f8652883cfc",
+   "execution_count": 21,
+   "id": "d70c40b9-5d5a-4795-b2a2-149c4a57d16e",
    "metadata": {},
    "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py:441: UserWarning: Examples will be cached but not all input components have example values. This may result in an exception being thrown by your function. If you do get an error while caching examples, make sure all of your inputs have example values for all of your examples or you provide default values for those particular parameters in your function.\n",
+      "  warnings.warn(\n",
+      "INFO:__main__:Using cached structure: ./7rpz.cif\n",
+      "INFO:__main__:Using cached structure: ./2iwi.cif\n",
+      "INFO:__main__:Using cached structure: ./2f6v.cif\n",
+      "INFO:httpx:HTTP Request: GET http://127.0.0.1:7862/gradio_api/startup-events \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on local URL:  http://127.0.0.1:7862\n",
+      "Caching examples at: '/home/frohlkin/Projects/LargeLanguageModels/Publication/test_webpage/.gradio/cached_examples/148'\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:httpx:HTTP Request: HEAD http://127.0.0.1:7862/ \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: GET https://api.gradio.app/pkg-version \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: GET https://api.gradio.app/v3/tunnel-request \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "* Running on local URL:  http://127.0.0.1:7860\n",
+      "* Running on public URL: https://de785d7cce806497e9.gradio.live\n",
       "\n",
-      "To create a public link, set `share=True` in `launch()`.\n"
+      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:httpx:HTTP Request: HEAD https://de785d7cce806497e9.gradio.live \"HTTP/1.1 200 OK\"\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>"
+       "<div><iframe src=\"https://de785d7cce806497e9.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"
@@ -644,79 +632,510 @@
      "output_type": "display_data"
     },
     {
-     "data": {
-      "text/plain": []
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Traceback (most recent call last):\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/routes.py\", line 990, in predict\n",
+      "    output = await route_utils.call_process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/route_utils.py\", line 322, in call_process_api\n",
+      "    output = await app.get_blocks().process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 2047, in process_api\n",
+      "    result = await self.call_function(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 1594, in call_function\n",
+      "    prediction = await anyio.to_thread.run_sync(  # type: ignore\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/to_thread.py\", line 56, in run_sync\n",
+      "    return await get_async_backend().run_sync_in_worker_thread(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 2405, in run_sync_in_worker_thread\n",
+      "    return await future\n",
+      "           ^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 914, in run\n",
+      "    result = context.run(func, *args)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/utils.py\", line 869, in wrapper\n",
+      "    response = f(*args, **kwargs)\n",
+      "               ^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 355, in load_example_with_output\n",
+      "    ) + self.load_from_cache(example_id)\n",
+      "        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 579, in load_from_cache\n",
+      "    output.append(component.read_from_flag(value_to_use))\n",
+      "                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/components/base.py\", line 366, in read_from_flag\n",
+      "    return self.data_model.from_json(json.loads(payload))\n",
+      "                                     ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/__init__.py\", line 346, in loads\n",
+      "    return _default_decoder.decode(s)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 337, in decode\n",
+      "    obj, end = self.raw_decode(s, idx=_w(s, 0).end())\n",
+      "               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 355, in raw_decode\n",
+      "    raise JSONDecodeError(\"Expecting value\", s, err.value) from None\n",
+      "json.decoder.JSONDecodeError: Expecting value: line 1 column 1 (char 0)\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/routes.py\", line 990, in predict\n",
+      "    output = await route_utils.call_process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/route_utils.py\", line 322, in call_process_api\n",
+      "    output = await app.get_blocks().process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 2047, in process_api\n",
+      "    result = await self.call_function(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 1594, in call_function\n",
+      "    prediction = await anyio.to_thread.run_sync(  # type: ignore\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/to_thread.py\", line 56, in run_sync\n",
+      "    return await get_async_backend().run_sync_in_worker_thread(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 2405, in run_sync_in_worker_thread\n",
+      "    return await future\n",
+      "           ^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 914, in run\n",
+      "    result = context.run(func, *args)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/utils.py\", line 869, in wrapper\n",
+      "    response = f(*args, **kwargs)\n",
+      "               ^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 355, in load_example_with_output\n",
+      "    ) + self.load_from_cache(example_id)\n",
+      "        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 579, in load_from_cache\n",
+      "    output.append(component.read_from_flag(value_to_use))\n",
+      "                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/components/base.py\", line 366, in read_from_flag\n",
+      "    return self.data_model.from_json(json.loads(payload))\n",
+      "                                     ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/__init__.py\", line 346, in loads\n",
+      "    return _default_decoder.decode(s)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 337, in decode\n",
+      "    obj, end = self.raw_decode(s, idx=_w(s, 0).end())\n",
+      "               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 355, in raw_decode\n",
+      "    raise JSONDecodeError(\"Expecting value\", s, err.value) from None\n",
+      "json.decoder.JSONDecodeError: Expecting value: line 1 column 1 (char 0)\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/routes.py\", line 990, in predict\n",
+      "    output = await route_utils.call_process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/route_utils.py\", line 322, in call_process_api\n",
+      "    output = await app.get_blocks().process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 2047, in process_api\n",
+      "    result = await self.call_function(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 1594, in call_function\n",
+      "    prediction = await anyio.to_thread.run_sync(  # type: ignore\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/to_thread.py\", line 56, in run_sync\n",
+      "    return await get_async_backend().run_sync_in_worker_thread(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 2405, in run_sync_in_worker_thread\n",
+      "    return await future\n",
+      "           ^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 914, in run\n",
+      "    result = context.run(func, *args)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/utils.py\", line 869, in wrapper\n",
+      "    response = f(*args, **kwargs)\n",
+      "               ^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 355, in load_example_with_output\n",
+      "    ) + self.load_from_cache(example_id)\n",
+      "        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 579, in load_from_cache\n",
+      "    output.append(component.read_from_flag(value_to_use))\n",
+      "                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/components/base.py\", line 366, in read_from_flag\n",
+      "    return self.data_model.from_json(json.loads(payload))\n",
+      "                                     ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/__init__.py\", line 346, in loads\n",
+      "    return _default_decoder.decode(s)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 337, in decode\n",
+      "    obj, end = self.raw_decode(s, idx=_w(s, 0).end())\n",
+      "               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 355, in raw_decode\n",
+      "    raise JSONDecodeError(\"Expecting value\", s, err.value) from None\n",
+      "json.decoder.JSONDecodeError: Expecting value: line 1 column 1 (char 0)\n"
+     ]
     }
    ],
    "source": [
+    "from datetime import datetime\n",
     "import gradio as gr\n",
     "import requests\n",
-    "from Bio.PDB import PDBParser\n",
+    "from Bio.PDB import PDBParser, MMCIFParser, PDBIO, Select, Structure\n",
+    "from Bio.PDB.Polypeptide import is_aa\n",
+    "from Bio.SeqUtils import seq1\n",
+    "from typing import Optional, Tuple, Dict, List\n",
     "import numpy as np\n",
     "import os\n",
     "from gradio_molecule3d import Molecule3D\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader\n",
+    "import re\n",
+    "import pandas as pd\n",
+    "import copy\n",
+    "from scipy.special import expit\n",
+    "import logging\n",
+    "import tempfile\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",
+    "# Set up logging\n",
+    "logging.basicConfig(level=logging.INFO)\n",
+    "logger = logging.getLogger(__name__)\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",
+    "class StructureError(Exception):\n",
+    "    \"\"\"Custom exception for structure-related errors\"\"\"\n",
+    "    pass\n",
+    "\n",
+    "def normalize_scores(scores: np.ndarray) -> np.ndarray:\n",
+    "    \"\"\"Normalize scores to range [0,1]\"\"\"\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: str) -> str:\n",
+    "    \"\"\"Read molecular structure file and return its content\"\"\"\n",
+    "    try:\n",
+    "        with open(pdb_path, 'r') as f:\n",
+    "            return f.read()\n",
+    "    except Exception as e:\n",
+    "        raise IOError(f\"Failed to read structure file: {e}\")\n",
+    "\n",
+    "def fetch_structure(pdb_id: str, output_dir: str = \".\") -> Optional[str]:\n",
+    "    \"\"\"Fetch structure file, trying multiple formats and sources\"\"\"\n",
+    "    try:\n",
+    "        # First try local cache\n",
+    "        for ext in ['.cif', '.pdb']:\n",
+    "            local_path = os.path.join(output_dir, f\"{pdb_id.lower()}{ext}\")\n",
+    "            if os.path.exists(local_path):\n",
+    "                logger.info(f\"Using cached structure: {local_path}\")\n",
+    "                return local_path\n",
+    "\n",
+    "        # Try different download sources\n",
+    "        sources = [\n",
+    "            f\"https://files.rcsb.org/download/{pdb_id.upper()}.cif\",\n",
+    "            f\"https://files.rcsb.org/download/{pdb_id.upper()}.pdb\",\n",
+    "            f\"https://files.rcsb.org/download/{pdb_id.lower()}.cif\",\n",
+    "            f\"https://files.rcsb.org/download/{pdb_id.lower()}.pdb\"\n",
+    "        ]\n",
+    "\n",
+    "        for url in sources:\n",
+    "            try:\n",
+    "                response = requests.get(url, timeout=10)\n",
+    "                if response.status_code == 200:\n",
+    "                    ext = '.cif' if 'cif' in url else '.pdb'\n",
+    "                    file_path = os.path.join(output_dir, f\"{pdb_id.lower()}{ext}\")\n",
+    "                    with open(file_path, 'wb') as f:\n",
+    "                        f.write(response.content)\n",
+    "                    logger.info(f\"Successfully downloaded: {url}\")\n",
+    "                    return file_path\n",
+    "            except Exception as e:\n",
+    "                logger.warning(f\"Failed to download from {url}: {e}\")\n",
+    "                continue\n",
+    "\n",
+    "        raise StructureError(f\"Failed to fetch structure for PDB ID: {pdb_id}\")\n",
+    "    except Exception as e:\n",
+    "        raise StructureError(f\"Error fetching structure: {e}\")\n",
+    "\n",
+    "def convert_cif_to_pdb(cif_path: str, output_dir: str = \".\") -> str:\n",
+    "    \"\"\"Convert CIF to PDB format with error handling\"\"\"\n",
+    "    try:\n",
+    "        pdb_path = os.path.join(output_dir, os.path.basename(cif_path).replace('.cif', '.pdb'))\n",
+    "        parser = MMCIFParser(QUIET=True)\n",
+    "        structure = parser.get_structure('protein', cif_path)\n",
+    "        io = PDBIO()\n",
+    "        io.set_structure(structure)\n",
+    "        io.save(pdb_path)\n",
     "        return pdb_path\n",
-    "    else:\n",
-    "        return None\n",
+    "    except Exception as e:\n",
+    "        raise StructureError(f\"Failed to convert CIF to PDB: {e}\")\n",
+    "\n",
+    "def find_valid_chain(structure: Structure.Structure) -> Optional[str]:\n",
+    "    \"\"\"Find the first valid protein chain in the structure\"\"\"\n",
+    "    for model in structure:\n",
+    "        for chain in model:\n",
+    "            protein_residues = [res for res in chain if is_aa(res)]\n",
+    "            if len(protein_residues) > 0:\n",
+    "                return chain.id\n",
+    "    return None\n",
+    "\n",
+    "def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list, protein_residues: list) -> str:\n",
+    "    \"\"\"Create PDB file with selected chain and prediction scores in B-factor column\"\"\"\n",
+    "    class ResidueSelector(Select):\n",
+    "        def __init__(self, chain_id, selected_residues, scores_dict):\n",
+    "            self.chain_id = chain_id\n",
+    "            self.selected_residues = selected_residues\n",
+    "            self.scores_dict = scores_dict\n",
+    "        \n",
+    "        def accept_chain(self, chain):\n",
+    "            return chain.id == self.chain_id\n",
+    "        \n",
+    "        def accept_residue(self, residue):\n",
+    "            return residue.id[1] in self.selected_residues\n",
+    "\n",
+    "        def accept_atom(self, atom):\n",
+    "            if atom.parent.id[1] in self.scores_dict:\n",
+    "                atom.bfactor = np.absolute(1-self.scores_dict[atom.parent.id[1]]) * 100\n",
+    "            return True\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",
+    "        parser = PDBParser(QUIET=True)\n",
+    "        structure = parser.get_structure('protein', input_pdb)\n",
+    "        output_pdb = f\"{os.path.splitext(input_pdb)[0]}_{chain_id}_predictions_scores.pdb\"\n",
+    "        scores_dict = {resi: score for resi, score in residue_scores}\n",
+    "        \n",
+    "        io = PDBIO()\n",
+    "        selector = ResidueSelector(chain_id, [res.id[1] for res in protein_residues], scores_dict)\n",
+    "        io.set_structure(structure[0])\n",
+    "        io.save(output_pdb, selector)\n",
+    "        \n",
+    "        return output_pdb\n",
+    "    except Exception as e:\n",
+    "        raise StructureError(f\"Failed to create chain-specific PDB: {e}\")\n",
+    "\n",
+    "def process_pdb(pdb_id_or_file: str, segment: str) -> Tuple[str, str, List[str]]:\n",
+    "    \"\"\"Process PDB/CIF file and generate visualizations and predictions\"\"\"\n",
+    "    try:\n",
+    "        # Handle input\n",
+    "        if pdb_id_or_file.endswith(('.pdb', '.cif')):\n",
+    "            pdb_path = pdb_id_or_file\n",
+    "            pdb_id = os.path.splitext(os.path.basename(pdb_path))[0]\n",
+    "        else:\n",
+    "            pdb_id = pdb_id_or_file\n",
+    "            pdb_path = fetch_structure(pdb_id)\n",
+    "\n",
+    "        if not pdb_path:\n",
+    "            raise StructureError(\"Failed to obtain structure file\")\n",
+    "\n",
+    "        # Parse structure\n",
+    "        parser = MMCIFParser(QUIET=True) if pdb_path.endswith('.cif') else PDBParser(QUIET=True)\n",
+    "        structure = parser.get_structure('protein', pdb_path)\n",
+    "\n",
+    "        # Handle chain selection\n",
+    "        if segment == 'auto' or not segment:\n",
+    "            segment = find_valid_chain(structure)\n",
+    "            if not segment:\n",
+    "                raise StructureError(\"No valid protein chains found in structure\")\n",
+    "        \n",
+    "        try:\n",
+    "            chain = structure[0][segment]\n",
+    "        except KeyError:\n",
+    "            valid_chain = find_valid_chain(structure)\n",
+    "            if valid_chain:\n",
+    "                chain = structure[0][valid_chain]\n",
+    "                segment = valid_chain\n",
+    "                logger.info(f\"Using alternative chain {segment}\")\n",
+    "            else:\n",
+    "                raise StructureError(f\"Invalid chain ID '{segment}'. Structure has no valid protein chains.\")\n",
+    "\n",
+    "        # Process chain\n",
+    "        protein_residues = [res for res in chain if is_aa(res)]\n",
+    "        if not protein_residues:\n",
+    "            raise StructureError(f\"No amino acid residues found in chain {segment}\")\n",
+    "\n",
+    "        sequence = \"\".join(seq1(res.resname) for res in protein_residues)\n",
+    "        sequence_id = [res.id[1] for res in protein_residues]\n",
+    "        \n",
+    "        # Generate predictions (currently random)\n",
+    "        scores = np.random.rand(len(sequence))\n",
+    "        normalized_scores = normalize_scores(scores)\n",
+    "        residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]\n",
+    "\n",
+    "        # Generate outputs\n",
+    "        result_str = generate_results_string(pdb_id, segment, protein_residues, normalized_scores, sequence)\n",
+    "        scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores, protein_residues)\n",
+    "        mol_vis = molecule(pdb_path, residue_scores, segment)\n",
+    "        pymol_commands = generate_pymol_commands(pdb_id, segment, residue_scores, pdb_path)\n",
+    "\n",
+    "        # Save results\n",
+    "        prediction_file = f\"{pdb_id}_binding_site_residues.txt\"\n",
+    "        with open(prediction_file, \"w\") as f:\n",
+    "            f.write(result_str)\n",
+    "\n",
+    "        return pymol_commands, mol_vis, [prediction_file, scored_pdb]\n",
+    "\n",
+    "    except StructureError as e:\n",
+    "        return str(e), None, None\n",
+    "    except Exception as e:\n",
+    "        return f\"An unexpected error occurred: {str(e)}\", None, None\n",
+    "\n",
+    "def generate_results_string(pdb_id: str, segment: str, protein_residues: list, \n",
+    "                          normalized_scores: np.ndarray, sequence: str) -> str:\n",
+    "    \"\"\"Generate formatted results string with predictions\"\"\"\n",
+    "    score_brackets = {\n",
+    "        \"0.0-0.2\": (0.0, 0.2),\n",
+    "        \"0.2-0.4\": (0.2, 0.4),\n",
+    "        \"0.4-0.6\": (0.4, 0.6),\n",
+    "        \"0.6-0.8\": (0.6, 0.8),\n",
+    "        \"0.8-1.0\": (0.8, 1.0)\n",
+    "    }\n",
     "    \n",
-    "    return result_str, molecule(pdb_path, random_scores), prediction_file\n",
+    "    residues_by_bracket = {bracket: [] for bracket in score_brackets}\n",
+    "    \n",
+    "    # Categorize residues\n",
+    "    for i, score in enumerate(normalized_scores):\n",
+    "        for bracket, (lower, upper) in score_brackets.items():\n",
+    "            if lower <= score < upper:\n",
+    "                residues_by_bracket[bracket].append(protein_residues[i])\n",
+    "                break\n",
+    "    \n",
+    "    # Format results\n",
+    "    current_time = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+    "    result_str = f\"\"\"Prediction Results\n",
+    "========================\n",
+    "PDB: {pdb_id}\n",
+    "Chain: {segment}\n",
+    "Date: {current_time}\n",
     "\n",
-    "def molecule(input_pdb, scores=None):\n",
-    "    mol = read_mol(input_pdb)  # Read PDB file content\n",
+    "Analysis by Score Brackets\n",
+    "========================\n",
+    "\"\"\"\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",
+    "    for bracket, residues in residues_by_bracket.items():\n",
+    "        if residues:  # Only show brackets with residues\n",
+    "            result_str += f\"\\nBracket {bracket}:\\n\"\n",
+    "            result_str += \"ResName ResNum Code Score\\n\"\n",
+    "            result_str += \"-\" * 30 + \"\\n\"\n",
+    "            result_str += \"\\n\".join([\n",
+    "                f\"{res.resname:6} {res.id[1]:6} {sequence[i]:4} {normalized_scores[i]:6.2f}\" \n",
+    "                for i, res in enumerate(protein_residues) if res in residues\n",
+    "            ])\n",
+    "            result_str += \"\\n\"\n",
+    "    \n",
+    "    return result_str\n",
+    "\n",
+    "def generate_pymol_commands(pdb_id: str, segment: str, residue_scores: list, pdb_path: str) -> str:\n",
+    "    \"\"\"Generate PyMOL visualization commands\"\"\"\n",
+    "    # Group residues by score ranges\n",
+    "    score_groups = {\n",
+    "        \"very_low\": [], \"low\": [], \"medium\": [], \"high\": [], \"very_high\": []\n",
+    "    }\n",
     "    \n",
+    "    for resi, score in residue_scores:\n",
+    "        if score <= 0.2:\n",
+    "            score_groups[\"very_low\"].append(str(resi))\n",
+    "        elif score <= 0.4:\n",
+    "            score_groups[\"low\"].append(str(resi))\n",
+    "        elif score <= 0.6:\n",
+    "            score_groups[\"medium\"].append(str(resi))\n",
+    "        elif score <= 0.8:\n",
+    "            score_groups[\"high\"].append(str(resi))\n",
+    "        else:\n",
+    "            score_groups[\"very_high\"].append(str(resi))\n",
+    "\n",
+    "    current_time = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+    "    commands = f\"\"\"# PyMOL Script for {pdb_id} Chain {segment}\n",
+    "# Generated: {current_time}\n",
+    "\n",
+    "# Load structure and set initial display\n",
+    "load {os.path.abspath(pdb_path)}, protein\n",
+    "bg_color white\n",
+    "hide everything\n",
+    "show cartoon, chain {segment}\n",
+    "color white, chain {segment}\n",
+    "\n",
+    "# Create selection groups by score\n",
+    "\"\"\"\n",
+    "    \n",
+    "    color_scheme = {\n",
+    "        \"very_low\": \"white\",\n",
+    "        \"low\": \"lightorange\",\n",
+    "        \"medium\": \"orange\",\n",
+    "        \"high\": \"orangered\",\n",
+    "        \"very_high\": \"red\"\n",
+    "    }\n",
+    "    \n",
+    "    for group, residues in score_groups.items():\n",
+    "        if residues:\n",
+    "            resi_str = \"+\".join(residues)\n",
+    "            commands += f\"\"\"\n",
+    "# {group.replace('_', ' ').title()} scoring residues\n",
+    "select {group}, chain {segment} and resi {resi_str}\n",
+    "show sticks, {group}\n",
+    "color {color_scheme[group]}, {group}\"\"\"\n",
+    "    \n",
+    "    commands += \"\"\"\n",
+    "\n",
+    "# Center and zoom\n",
+    "center chain {}\n",
+    "zoom chain {}\n",
+    "\"\"\"\n",
+    "\n",
+    "    return commands\n",
+    "\n",
+    "def molecule(input_pdb: str, residue_scores: list = None, segment: str = 'A') -> str:\n",
+    "    \"\"\"Generate interactive 3D molecule visualization\"\"\"\n",
+    "    try:\n",
+    "        mol = read_mol(input_pdb)\n",
+    "    except Exception as e:\n",
+    "        return f'<div class=\"error\">Error loading structure: {str(e)}</div>'\n",
+    "\n",
+    "    # Prepare residue groups for visualization\n",
+    "    vis_groups = {\n",
+    "        \"class1\": [],  # 0.0-0.2\n",
+    "        \"class2\": [],  # 0.2-0.4\n",
+    "        \"class3\": [],  # 0.4-0.6\n",
+    "        \"class4\": [],  # 0.6-0.8\n",
+    "        \"class5\": []   # 0.8-1.0\n",
+    "    }\n",
+    "\n",
+    "    if residue_scores:\n",
+    "        for resi, score in residue_scores:\n",
+    "            if score <= 0.2:\n",
+    "                vis_groups[\"class1\"].append(resi)\n",
+    "            elif score <= 0.4:\n",
+    "                vis_groups[\"class2\"].append(resi)\n",
+    "            elif score <= 0.6:\n",
+    "                vis_groups[\"class3\"].append(resi)\n",
+    "            elif score <= 0.8:\n",
+    "                vis_groups[\"class4\"].append(resi)\n",
+    "            else:\n",
+    "                vis_groups[\"class5\"].append(resi)\n",
+    "\n",
+    "    # Generate visualization script\n",
+    "    vis_script = f\"\"\"\n",
+    "        // Base model setup\n",
+    "        let chainModel = viewer.addModel(pdb, \"pdb\");\n",
+    "        chainModel.setStyle({{}}, {{}});\n",
+    "        chainModel.setStyle(\n",
+    "            {{\"chain\": \"{segment}\"}}, \n",
+    "            {{\"cartoon\": {{\"color\": \"white\"}}}}\n",
+    "        );\n",
+    "    \"\"\"\n",
+    "\n",
+    "    # Color schemes for different score ranges\n",
+    "    color_schemes = {\n",
+    "        \"class1\": {\"color\": \"0xFFFFFF\", \"opacity\": 0.5},  # White\n",
+    "        \"class2\": {\"color\": \"0xFFD580\", \"opacity\": 0.7},  # Light orange\n",
+    "        \"class3\": {\"color\": \"0xFFA500\", \"opacity\": 1.0},  # Orange\n",
+    "        \"class4\": {\"color\": \"0xFF4500\", \"opacity\": 1.0},  # Orange red\n",
+    "        \"class5\": {\"color\": \"0xFF0000\", \"opacity\": 1.0}   # Red\n",
+    "    }\n",
+    "\n",
+    "    # Add visualization for each group\n",
+    "    for group, residues in vis_groups.items():\n",
+    "        if residues:\n",
+    "            color_scheme = color_schemes[group]\n",
+    "            vis_script += f\"\"\"\n",
+    "        let {group}Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        {group}Model.setStyle({{}}, {{}});\n",
+    "        {group}Model.setStyle(\n",
+    "            {{\"chain\": \"{segment}\", \"resi\": [{\", \".join(map(str, residues))}]}},\n",
+    "            {{\"stick\": {{\"color\": \"{color_scheme[\"color\"]}\", \"opacity\": {color_scheme[\"opacity\"]}}}}}\n",
+    "        );\n",
+    "        \"\"\"\n",
+    "\n",
+    "    # Generate full HTML with enhanced controls and information\n",
     "    html_content = f\"\"\"\n",
     "    <!DOCTYPE html>\n",
     "    <html>\n",
@@ -728,95 +1147,343 @@
     "            height: 700px;\n",
     "            position: relative;\n",
     "        }}\n",
+    "        .controls {{\n",
+    "            position: absolute;\n",
+    "            top: 10px;\n",
+    "            left: 10px;\n",
+    "            background: rgba(255, 255, 255, 0.8);\n",
+    "            padding: 10px;\n",
+    "            border-radius: 5px;\n",
+    "            z-index: 1000;\n",
+    "        }}\n",
+    "        .legend {{\n",
+    "            position: absolute;\n",
+    "            bottom: 10px;\n",
+    "            right: 10px;\n",
+    "            background: rgba(255, 255, 255, 0.8);\n",
+    "            padding: 10px;\n",
+    "            border-radius: 5px;\n",
+    "            z-index: 1000;\n",
+    "        }}\n",
+    "        .error {{\n",
+    "            color: red;\n",
+    "            padding: 20px;\n",
+    "            text-align: center;\n",
+    "            font-weight: bold;\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",
+    "        <div id=\"container\" class=\"mol-container\">\n",
+    "            <div class=\"controls\">\n",
+    "                <button onclick=\"toggleStyle('cartoon')\">Toggle Cartoon</button>\n",
+    "                <button onclick=\"toggleStyle('stick')\">Toggle Sticks</button>\n",
+    "                <button onclick=\"resetView()\">Reset View</button>\n",
+    "                <button onclick=\"toggleSpin()\">Toggle Spin</button>\n",
+    "            </div>\n",
+    "            <div class=\"legend\">\n",
+    "                <div><span style=\"color: #FF0000\">■</span> Very High (0.8-1.0)</div>\n",
+    "                <div><span style=\"color: #FF4500\">■</span> High (0.6-0.8)</div>\n",
+    "                <div><span style=\"color: #FFA500\">■</span> Medium (0.4-0.6)</div>\n",
+    "                <div><span style=\"color: #FFD580\">■</span> Low (0.2-0.4)</div>\n",
+    "                <div><span style=\"color: #FFFFFF\">■</span> Very Low (0.0-0.2)</div>\n",
+    "            </div>\n",
+    "        </div>\n",
     "        <script>\n",
-    "            let pdb = `{mol}`; // Use template literal to properly escape PDB content\n",
+    "            let pdb = `{mol}`;\n",
+    "            let viewer;\n",
+    "            let isSpinning = false;\n",
+    "\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 = $3Dmol.createViewer(element, config);\n",
     "                \n",
-    "                // Set cartoon representation with white carbon color scheme\n",
-    "                viewer.getModel(0).setStyle({{}}, {{ cartoon: {{ colorscheme:\"whiteCarbon\" }} }});\n",
+    "                {vis_script}\n",
     "                \n",
-    "                {high_score_script}\n",
+    "                // Enhanced 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: 18,\n",
+    "                                    padding: 4\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",
+    "\n",
+    "            function toggleStyle(style) {{\n",
+    "                let elements = viewer.selectedAtoms({{}});\n",
+    "                let currentStyle = elements.style[style];\n",
+    "                elements.style[style] = !currentStyle;\n",
+    "                viewer.render();\n",
+    "            }}\n",
+    "\n",
+    "            function resetView() {{\n",
+    "                viewer.zoomTo();\n",
+    "                viewer.render();\n",
+    "            }}\n",
+    "\n",
+    "            function toggleSpin() {{\n",
+    "                isSpinning = !isSpinning;\n",
+    "                viewer.spin(isSpinning);\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",
+    "def create_ui():\n",
+    "    with gr.Blocks(title=\"Protein Binding Site Prediction\", theme=gr.themes.Base()) as demo:\n",
+    "        gr.Markdown(\"\"\"\n",
+    "        # Protein Binding Site Prediction\n",
+    "        \n",
+    "        This tool helps you visualize and analyze potential binding sites in protein structures.\n",
+    "        You can either:\n",
+    "        1. Enter a PDB ID (e.g., \"4BDU\")\n",
+    "        2. Upload your own PDB/CIF file\n",
+    "        \n",
+    "        The tool will analyze the structure and show predictions using a color gradient from white (low probability) to red (high probability).\n",
+    "        \"\"\")\n",
+    "        \n",
+    "        with gr.Row():\n",
+    "            with gr.Column(scale=2):\n",
+    "                # Input components\n",
+    "                mode = gr.Radio(\n",
+    "                    choices=[\"PDB ID\", \"Upload File\"],\n",
+    "                    value=\"PDB ID\",\n",
+    "                    label=\"Input Mode\",\n",
+    "                    info=\"Choose whether to input a PDB ID or upload a PDB/CIF file\"\n",
+    "                )\n",
+    "                \n",
+    "                with gr.Group():\n",
+    "                    pdb_input = gr.Textbox(\n",
+    "                        value=\"4BDU\",\n",
+    "                        label=\"PDB ID\",\n",
+    "                        placeholder=\"Enter PDB ID (e.g., 4BDU)\",\n",
+    "                        info=\"Enter a valid PDB ID from the Protein Data Bank\"\n",
+    "                    )\n",
+    "                    pdb_file = gr.File(\n",
+    "                        label=\"Upload PDB/CIF File\",\n",
+    "                        file_types=[\".pdb\", \".cif\"],\n",
+    "                        visible=False\n",
+    "                    )\n",
+    "                \n",
+    "                segment_input = gr.Textbox(\n",
+    "                    value=\"A\",\n",
+    "                    label=\"Chain ID\",\n",
+    "                    placeholder=\"Enter Chain ID or leave empty for automatic selection\",\n",
+    "                    info=\"Specify which protein chain to analyze, or leave empty for automatic selection\"\n",
+    "                )\n",
     "\n",
-    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=reps)\n",
+    "            with gr.Column(scale=1):\n",
+    "                visualize_btn = gr.Button(\"Visualize Structure\", variant=\"primary\")\n",
+    "                prediction_btn = gr.Button(\"Predict Binding Site\", variant=\"secondary\")\n",
+    "                \n",
+    "                gr.Markdown(\"\"\"\n",
+    "                ### Color Legend\n",
+    "                - White: Very Low (0.0-0.2)\n",
+    "                - Light Orange: Low (0.2-0.4)\n",
+    "                - Orange: Medium (0.4-0.6)\n",
+    "                - Orange Red: High (0.6-0.8)\n",
+    "                - Red: Very High (0.8-1.0)\n",
+    "                \"\"\")\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",
+    "        with gr.Tab(\"3D Visualization\"):\n",
+    "            molecule_output = gr.HTML(label=\"Interactive 3D Structure\")\n",
+    "            \n",
+    "        with gr.Tab(\"Analysis Results\"):\n",
+    "            predictions_output = gr.Textbox(\n",
+    "                label=\"PyMOL Visualization Commands\",\n",
+    "                info=\"Copy these commands into PyMOL to recreate the visualization\"\n",
+    "            )\n",
+    "            download_output = gr.File(\n",
+    "                label=\"Download Results\",\n",
+    "                file_count=\"multiple\"\n",
+    "            )\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",
+    "        # Error message container\n",
+    "        error_output = gr.Markdown(visible=False)\n",
+    "\n",
+    "        # Mode change handler\n",
+    "        def toggle_mode(selected_mode):\n",
+    "            return {\n",
+    "                pdb_input: gr.update(visible=selected_mode == \"PDB ID\"),\n",
+    "                pdb_file: gr.update(visible=selected_mode == \"Upload File\")\n",
+    "            }\n",
+    "\n",
+    "        mode.change(\n",
+    "            toggle_mode,\n",
+    "            inputs=[mode],\n",
+    "            outputs=[pdb_input, pdb_file]\n",
+    "        )\n",
+    "\n",
+    "        # Process handlers\n",
+    "        def handle_visualization(mode, pdb_id, pdb_file):\n",
+    "            try:\n",
+    "                result = fetch_interface(mode, pdb_id, pdb_file)\n",
+    "                if isinstance(result, str) and result.startswith(\"Error\"):\n",
+    "                    return None, gr.update(visible=True, value=f\"```\\n{result}\\n```\")\n",
+    "                return result, gr.update(visible=False)\n",
+    "            except Exception as e:\n",
+    "                return None, gr.update(visible=True, value=f\"```\\nError: {str(e)}\\n```\")\n",
+    "\n",
+    "        def handle_prediction(mode, pdb_id, pdb_file, chain_id):\n",
+    "            try:\n",
+    "                predictions, vis, downloads = process_interface(mode, pdb_id, pdb_file, chain_id)\n",
+    "                if isinstance(predictions, str) and \"Error\" in predictions:\n",
+    "                    return (\n",
+    "                        None,\n",
+    "                        None,\n",
+    "                        None,\n",
+    "                        gr.update(visible=True, value=f\"```\\n{predictions}\\n```\")\n",
+    "                    )\n",
+    "                return (\n",
+    "                    predictions,\n",
+    "                    vis,\n",
+    "                    downloads,\n",
+    "                    gr.update(visible=False)\n",
+    "                )\n",
+    "            except Exception as e:\n",
+    "                error_msg = f\"\"\"Error processing structure:\n",
+    "```\n",
+    "{str(e)}\n",
+    "\n",
+    "Troubleshooting tips:\n",
+    "1. Check if the PDB ID is valid\n",
+    "2. Ensure the Chain ID exists in the structure\n",
+    "3. Try leaving Chain ID empty for automatic selection\n",
+    "4. If uploading a file, ensure it's a valid PDB/CIF format\n",
+    "```\"\"\"\n",
+    "                return None, None, None, gr.update(visible=True, value=error_msg)\n",
+    "\n",
+    "        def fetch_interface(mode, pdb_id, pdb_file):\n",
+    "            try:\n",
+    "                if mode == \"PDB ID\":\n",
+    "                    if not pdb_id or len(pdb_id.strip()) != 4:\n",
+    "                        raise ValueError(\"Please enter a valid 4-character PDB ID\")\n",
+    "                    return fetch_pdb(pdb_id.strip())\n",
+    "                elif mode == \"Upload File\":\n",
+    "                    if not pdb_file:\n",
+    "                        raise ValueError(\"Please upload a PDB or CIF file\")\n",
+    "                    _, ext = os.path.splitext(pdb_file.name)\n",
+    "                    if ext.lower() not in ['.pdb', '.cif']:\n",
+    "                        raise ValueError(\"Only .pdb and .cif files are supported\")\n",
+    "                    \n",
+    "                    # Create temp directory for file handling\n",
+    "                    with tempfile.TemporaryDirectory() as temp_dir:\n",
+    "                        temp_path = os.path.join(temp_dir, os.path.basename(pdb_file.name))\n",
+    "                        with open(temp_path, 'wb') as f:\n",
+    "                            f.write(pdb_file.read())\n",
+    "                        \n",
+    "                        if ext.lower() == '.cif':\n",
+    "                            return convert_cif_to_pdb(temp_path)\n",
+    "                        return temp_path\n",
+    "                else:\n",
+    "                    raise ValueError(\"Invalid mode selected\")\n",
+    "            except Exception as e:\n",
+    "                return f\"Error: {str(e)}\"\n",
+    "\n",
+    "        # Connect event handlers\n",
+    "        visualize_btn.click(\n",
+    "            handle_visualization,\n",
+    "            inputs=[mode, pdb_input, pdb_file],\n",
+    "            outputs=[molecule_output, error_output]\n",
+    "        )\n",
     "\n",
-    "demo.launch()"
+    "        prediction_btn.click(\n",
+    "            handle_prediction,\n",
+    "            inputs=[mode, pdb_input, pdb_file, segment_input],\n",
+    "            outputs=[predictions_output, molecule_output, download_output, error_output]\n",
+    "        )\n",
+    "\n",
+    "        # Add examples\n",
+    "        gr.Examples(\n",
+    "            examples=[\n",
+    "                [\"PDB ID\", \"7RPZ\", None, \"A\"],\n",
+    "                [\"PDB ID\", \"2IWI\", None, \"B\"],\n",
+    "                [\"PDB ID\", \"2F6V\", None, \"A\"]\n",
+    "            ],\n",
+    "            inputs=[mode, pdb_input, pdb_file, segment_input],\n",
+    "            outputs=[predictions_output, molecule_output, download_output, error_output],\n",
+    "            fn=handle_prediction,\n",
+    "            cache_examples=True\n",
+    "        )\n",
+    "\n",
+    "        # Add documentation\n",
+    "        gr.Markdown(\"\"\"\n",
+    "        ## Usage Instructions\n",
+    "        \n",
+    "        1. **Input Structure:**\n",
+    "           - Enter a PDB ID (e.g., \"4BDU\") or upload your own structure file\n",
+    "           - The tool supports both PDB (.pdb) and mmCIF (.cif) formats\n",
+    "        \n",
+    "        2. **Select Chain:**\n",
+    "           - Enter a specific chain ID (e.g., \"A\")\n",
+    "           - Leave empty for automatic selection of the first valid protein chain\n",
+    "        \n",
+    "        3. **Analyze:**\n",
+    "           - Click \"Visualize Structure\" to view the 3D structure\n",
+    "           - Click \"Predict Binding Site\" to perform binding site analysis\n",
+    "        \n",
+    "        4. **Results:**\n",
+    "           - Interactive 3D visualization with color-coded predictions\n",
+    "           - PyMOL commands for external visualization\n",
+    "           - Downloadable results files\n",
+    "        \n",
+    "        ## Troubleshooting\n",
+    "        \n",
+    "        If you encounter issues:\n",
+    "        1. Ensure your PDB ID is valid and exists in the PDB database\n",
+    "        2. Check that your uploaded file is a valid PDB/CIF format\n",
+    "        3. Try automatic chain selection if your specified chain isn't found\n",
+    "        4. Clear your browser cache if visualizations don't load\n",
+    "        \"\"\")\n",
+    "\n",
+    "    return demo\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    demo = create_ui()\n",
+    "    demo.launch(share=True)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "14605615-8610-4d9e-841b-db7618cde844",
+   "id": "9125d1c4-e2ae-4e40-ba36-7ae944512b8e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "85c0728a-a15b-4118-b920-5f55a2f5f79a",
    "metadata": {},
    "outputs": [],
    "source": []
@@ -838,7 +1505,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.7"
+   "version": "3.12.2"
   }
  },
  "nbformat": 4,

app.py

@@ -139,30 +139,6 @@ def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: lis
     
     return output_pdb
 
-def calculate_geometric_center(pdb_path: str, high_score_residues: list, chain_id: str):
-    """
-    Calculate the geometric center of high-scoring residues
-    """
-    parser = PDBParser(QUIET=True)
-    structure = parser.get_structure('protein', pdb_path)
-    
-    # Collect coordinates of CA atoms from high-scoring residues
-    coords = []
-    for model in structure:
-        for chain in model:
-            if chain.id == chain_id:
-                for residue in chain:
-                    if residue.id[1] in high_score_residues:
-                        if 'CA' in residue:  # Use alpha carbon as representative
-                            ca_atom = residue['CA']
-                            coords.append(ca_atom.coord)
-    
-    # Calculate geometric center
-    if coords:
-        center = np.mean(coords, axis=0)
-        return center
-    return None
-
 def process_pdb(pdb_id_or_file, segment):
     # Determine if input is a PDB ID or file path
     if pdb_id_or_file.endswith('.pdb'):
@@ -194,7 +170,11 @@ def process_pdb(pdb_id_or_file, segment):
     protein_residues = [res for res in chain if is_aa(res)]
     sequence = "".join(seq1(res.resname) for res in protein_residues)
     sequence_id = [res.id[1] for res in protein_residues]
-    
+
+    visualized_sequence = "".join(seq1(res.resname) for res in protein_residues)
+    if sequence != visualized_sequence:
+        raise ValueError("The visualized sequence does not match the prediction sequence")
+     
     input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
     with torch.no_grad():
         outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()
@@ -300,7 +280,6 @@ def molecule(input_pdb, residue_scores=None, segment='A'):
         class4_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.8]
         class5_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 1.0]
         
-        
         high_score_script = """
         // Load the original model and apply white cartoon style
         let chainModel = viewer.addModel(pdb, "pdb");
@@ -430,7 +409,19 @@ def molecule(input_pdb, residue_scores=None, segment='A'):
     return f'<iframe width="100%" height="700" srcdoc="{html_content.replace(chr(34), "&quot;").replace(chr(39), "&#39;")}"></iframe>'
 
 # Gradio UI
-with gr.Blocks() as demo:
+with gr.Blocks(css="""
+    /* Customize Gradio button colors */
+    #visualize-btn, #predict-btn {
+        background-color: #FF7300; /* Deep orange */
+        color: white;
+        border-radius: 5px;
+        padding: 10px;
+        font-weight: bold;
+    }
+    #visualize-btn:hover, #predict-btn:hover {
+        background-color: #CC5C00; /* Darkened orange on hover */
+    }
+""") as demo:
     gr.Markdown("# Protein Binding Site Prediction")
     
     # Mode selection
@@ -442,9 +433,9 @@ with gr.Blocks() as demo:
     )
 
     # Input components based on mode
-    pdb_input = gr.Textbox(value="4BDU", label="PDB ID", placeholder="Enter PDB ID here...")
+    pdb_input = gr.Textbox(value="2F6V", label="PDB ID", placeholder="Enter PDB ID here...")
     pdb_file = gr.File(label="Upload PDB/CIF File", visible=False)
-    visualize_btn = gr.Button("Visualize Structure")
+    visualize_btn = gr.Button("Visualize Structure", elem_id="visualize-btn")
 
     molecule_output2 = Molecule3D(label="Protein Structure", reps=[
         {
@@ -458,8 +449,9 @@ with gr.Blocks() as demo:
     ])
 
     with gr.Row():
-        segment_input = gr.Textbox(value="A", label="Chain ID", placeholder="Enter Chain ID here...")
-        prediction_btn = gr.Button("Predict Binding Site")
+        segment_input = gr.Textbox(value="A", label="Chain ID (protein)", placeholder="Enter Chain ID here...",
+        info="Choose in which chain to predict binding sites.")
+        prediction_btn = gr.Button("Predict Binding Site", elem_id="predict-btn")
 
     molecule_output = gr.HTML(label="Protein Structure")
     explanation_vis = gr.Markdown("""
@@ -533,7 +525,7 @@ with gr.Blocks() as demo:
         examples=[
             ["7RPZ", "A"],
             ["2IWI", "B"],
-            ["2F6V", "A"]
+            ["7LCJ", "R"]
         ],
         inputs=[pdb_input, segment_input],
         outputs=[predictions_output, molecule_output, download_output]

app.py.backup

@@ -0,0 +1,542 @@
+from datetime import datetime
+import gradio as gr
+import requests
+from Bio.PDB import PDBParser, MMCIFParser, PDBIO, Select
+from Bio.PDB.Polypeptide import is_aa
+from Bio.SeqUtils import seq1
+from typing import Optional, Tuple
+import numpy as np
+import os
+from gradio_molecule3d import Molecule3D
+
+from model_loader import load_model
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+import re
+import pandas as pd
+import copy
+
+import transformers
+from transformers import AutoTokenizer, DataCollatorForTokenClassification
+
+from datasets import Dataset
+
+from scipy.special import expit
+
+
+# Load model and move to device
+checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
+max_length = 1500
+model, tokenizer = load_model(checkpoint, max_length)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model.to(device)
+model.eval()
+
+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_structure(pdb_id: str, output_dir: str = ".") -> Optional[str]:
+    """
+    Fetch the structure file for a given PDB ID. Prioritizes CIF files.
+    If a structure file already exists locally, it uses that.
+    """
+    file_path = download_structure(pdb_id, output_dir)
+    if file_path:
+        return file_path
+    else:
+        return None
+
+def download_structure(pdb_id: str, output_dir: str) -> Optional[str]:
+    """
+    Attempt to download the structure file in CIF or PDB format.
+    Returns the path to the downloaded file, or None if download fails.
+    """
+    for ext in ['.cif', '.pdb']:
+        file_path = os.path.join(output_dir, f"{pdb_id}{ext}")
+        if os.path.exists(file_path):
+            return file_path
+        url = f"https://files.rcsb.org/download/{pdb_id}{ext}"
+        try:
+            response = requests.get(url, timeout=10)
+            if response.status_code == 200:
+                with open(file_path, 'wb') as f:
+                    f.write(response.content)
+                return file_path
+        except Exception as e:
+            print(f"Download error for {pdb_id}{ext}: {e}")
+    return None
+
+def convert_cif_to_pdb(cif_path: str, output_dir: str = ".") -> str:
+    """
+    Convert a CIF file to PDB format using BioPython and return the PDB file path.
+    """
+    pdb_path = os.path.join(output_dir, os.path.basename(cif_path).replace('.cif', '.pdb'))
+    parser = MMCIFParser(QUIET=True)
+    structure = parser.get_structure('protein', cif_path)
+    io = PDBIO()
+    io.set_structure(structure)
+    io.save(pdb_path)
+    return pdb_path
+
+def fetch_pdb(pdb_id):
+    pdb_path = fetch_structure(pdb_id)
+    if not pdb_path:
+        return None
+    _, ext = os.path.splitext(pdb_path)
+    if ext == '.cif':
+        pdb_path = convert_cif_to_pdb(pdb_path)
+    return pdb_path
+
+def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list, protein_residues: list) -> str:
+    """
+    Create a PDB file with only the selected chain and residues, replacing B-factor with prediction scores
+    """
+    # Read the original PDB file
+    parser = PDBParser(QUIET=True)
+    structure = parser.get_structure('protein', input_pdb)
+    
+    # Prepare a new structure with only the specified chain and selected residues
+    output_pdb = f"{os.path.splitext(input_pdb)[0]}_{chain_id}_predictions_scores.pdb"
+    
+    # Create scores dictionary for easy lookup
+    scores_dict = {resi: score for resi, score in residue_scores}
+
+    # Create a custom Select class
+    class ResidueSelector(Select):
+        def __init__(self, chain_id, selected_residues, scores_dict):
+            self.chain_id = chain_id
+            self.selected_residues = selected_residues
+            self.scores_dict = scores_dict
+        
+        def accept_chain(self, chain):
+            return chain.id == self.chain_id
+        
+        def accept_residue(self, residue):
+            return residue.id[1] in self.selected_residues
+
+        def accept_atom(self, atom):
+            if atom.parent.id[1] in self.scores_dict:
+                atom.bfactor = np.absolute(1-self.scores_dict[atom.parent.id[1]]) * 100
+            return True
+
+    # Prepare output PDB with selected chain and residues, modified B-factors
+    io = PDBIO()
+    selector = ResidueSelector(chain_id, [res.id[1] for res in protein_residues], scores_dict)
+    
+    io.set_structure(structure[0])
+    io.save(output_pdb, selector)
+    
+    return output_pdb
+
+def calculate_geometric_center(pdb_path: str, high_score_residues: list, chain_id: str):
+    """
+    Calculate the geometric center of high-scoring residues
+    """
+    parser = PDBParser(QUIET=True)
+    structure = parser.get_structure('protein', pdb_path)
+    
+    # Collect coordinates of CA atoms from high-scoring residues
+    coords = []
+    for model in structure:
+        for chain in model:
+            if chain.id == chain_id:
+                for residue in chain:
+                    if residue.id[1] in high_score_residues:
+                        if 'CA' in residue:  # Use alpha carbon as representative
+                            ca_atom = residue['CA']
+                            coords.append(ca_atom.coord)
+    
+    # Calculate geometric center
+    if coords:
+        center = np.mean(coords, axis=0)
+        return center
+    return None
+
+def process_pdb(pdb_id_or_file, segment):
+    # Determine if input is a PDB ID or file path
+    if pdb_id_or_file.endswith('.pdb'):
+        pdb_path = pdb_id_or_file
+        pdb_id = os.path.splitext(os.path.basename(pdb_path))[0]
+    else:
+        pdb_id = pdb_id_or_file
+        pdb_path = fetch_pdb(pdb_id)
+    
+    if not pdb_path:
+        return "Failed to fetch PDB file", None, None
+    
+    # Determine the file format and choose the appropriate parser
+    _, ext = os.path.splitext(pdb_path)
+    parser = MMCIFParser(QUIET=True) if ext == '.cif' else PDBParser(QUIET=True)
+    
+    try:
+        # Parse the structure file
+        structure = parser.get_structure('protein', pdb_path)
+    except Exception as e:
+        return f"Error parsing structure file: {e}", None, None
+    
+    # Extract the specified chain
+    try:
+        chain = structure[0][segment]
+    except KeyError:
+        return "Invalid Chain ID", None, None
+    
+    protein_residues = [res for res in chain if is_aa(res)]
+    sequence = "".join(seq1(res.resname) for res in protein_residues)
+    sequence_id = [res.id[1] for res in protein_residues]
+    
+    input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
+    with torch.no_grad():
+        outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()
+    
+    # Calculate scores and normalize them
+    scores = expit(outputs[:, 1] - outputs[:, 0])
+    
+    normalized_scores = normalize_scores(scores)
+    
+    # Zip residues with scores to track the residue ID and score
+    residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]
+
+    
+    # Define the score brackets
+    score_brackets = {
+        "0.0-0.2": (0.0, 0.2),
+        "0.2-0.4": (0.2, 0.4),
+        "0.4-0.6": (0.4, 0.6),
+        "0.6-0.8": (0.6, 0.8),
+        "0.8-1.0": (0.8, 1.0)
+    }
+    
+    # Initialize a dictionary to store residues by bracket
+    residues_by_bracket = {bracket: [] for bracket in score_brackets}
+    
+    # Categorize residues into brackets
+    for resi, score in residue_scores:
+        for bracket, (lower, upper) in score_brackets.items():
+            if lower <= score < upper:
+                residues_by_bracket[bracket].append(resi)
+                break
+    
+    # Preparing the result string
+    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    result_str = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\n\n"
+    result_str += "Residues by Score Brackets:\n\n"
+    
+    # Add residues for each bracket
+    for bracket, residues in residues_by_bracket.items():
+        result_str += f"Bracket {bracket}:\n"
+        result_str += "Columns: Residue Name, Residue Number, One-letter Code, Normalized Score\n"
+        result_str += "\n".join([
+            f"{res.resname} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}" 
+            for i, res in enumerate(protein_residues) if res.id[1] in residues
+        ])
+        result_str += "\n\n"
+
+    # Create chain-specific PDB with scores in B-factor
+    scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores, protein_residues)
+
+    # Molecule visualization with updated script with color mapping
+    mol_vis = molecule(pdb_path, residue_scores, segment)#, color_map)
+
+    # Improved PyMOL command suggestions
+    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    pymol_commands = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\n\n"
+    
+    pymol_commands += f"""
+    # PyMOL Visualization Commands
+    load {os.path.abspath(pdb_path)}, protein
+    hide everything, all
+    show cartoon, chain {segment}
+    color white, chain {segment}
+    """
+    
+    # Define colors for each score bracket
+    bracket_colors = {
+        "0.0-0.2": "white",
+        "0.2-0.4": "lightorange",
+        "0.4-0.6": "orange",
+        "0.6-0.8": "orangered",
+        "0.8-1.0": "red"
+    }
+    
+    # Add PyMOL commands for each score bracket
+    for bracket, residues in residues_by_bracket.items():
+        if residues:  # Only add commands if there are residues in this bracket
+            color = bracket_colors[bracket]
+            resi_list = '+'.join(map(str, residues))
+            pymol_commands += f"""
+    select bracket_{bracket.replace('.', '').replace('-', '_')}, resi {resi_list} and chain {segment}
+    show sticks, bracket_{bracket.replace('.', '').replace('-', '_')}
+    color {color}, bracket_{bracket.replace('.', '').replace('-', '_')}
+    """
+    # Create prediction and scored PDB files
+    prediction_file = f"{pdb_id}_binding_site_residues.txt"
+    with open(prediction_file, "w") as f:
+        f.write(result_str)
+    
+    return pymol_commands, mol_vis, [prediction_file,scored_pdb]
+
+def molecule(input_pdb, residue_scores=None, segment='A'):
+    # More granular scoring for visualization
+    mol = read_mol(input_pdb)  # Read PDB file content
+
+    # Prepare high-scoring residues script if scores are provided
+    high_score_script = ""
+    if residue_scores is not None:
+        # Filter residues based on their scores
+        class1_score_residues = [resi for resi, score in residue_scores if 0.0 < score <= 0.2]
+        class2_score_residues = [resi for resi, score in residue_scores if 0.2 < score <= 0.4]
+        class3_score_residues = [resi for resi, score in residue_scores if 0.4 < score <= 0.6]
+        class4_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.8]
+        class5_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 1.0]
+        
+        
+        high_score_script = """
+        // Load the original model and apply white cartoon style
+        let chainModel = viewer.addModel(pdb, "pdb");
+        chainModel.setStyle({}, {});
+        chainModel.setStyle(
+            {"chain": "%s"}, 
+            {"cartoon": {"color": "white"}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class1Model = viewer.addModel(pdb, "pdb");
+        class1Model.setStyle({}, {});
+        class1Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "0xFFFFFF", "opacity": 0.5}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class2Model = viewer.addModel(pdb, "pdb");
+        class2Model.setStyle({}, {});
+        class2Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "0xFFD580", "opacity": 0.7}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class3Model = viewer.addModel(pdb, "pdb");
+        class3Model.setStyle({}, {});
+        class3Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "0xFFA500", "opacity": 1}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class4Model = viewer.addModel(pdb, "pdb");
+        class4Model.setStyle({}, {});
+        class4Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "0xFF4500", "opacity": 1}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class5Model = viewer.addModel(pdb, "pdb");
+        class5Model.setStyle({}, {});
+        class5Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "0xFF0000", "alpha": 1}}
+        );
+
+        """ % (
+            segment,
+            segment,
+            ", ".join(str(resi) for resi in class1_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class2_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class3_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class4_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class5_score_residues)
+        )
+    
+    # Generate the full HTML content
+    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);
+                
+                {high_score_script}
+                
+                // Add hover functionality
+                viewer.setHoverable(
+                    {{}}, 
+                    true, 
+                    function(atom, viewer, event, container) {{
+                        if (!atom.label) {{
+                            atom.label = viewer.addLabel(
+                                atom.resn + ":" +atom.resi + ":" + atom.atom, 
+                                {{
+                                    position: atom, 
+                                    backgroundColor: 'mintcream', 
+                                    fontColor: 'black',
+                                    fontSize: 18,
+                                    padding: 4
+                                }}
+                            );
+                        }}
+                    }},
+                    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>'
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# Protein Binding Site Prediction")
+    
+    # Mode selection
+    mode = gr.Radio(
+        choices=["PDB ID", "Upload File"],
+        value="PDB ID",
+        label="Input Mode",
+        info="Choose whether to input a PDB ID or upload a PDB/CIF file."
+    )
+
+    # Input components based on mode
+    pdb_input = gr.Textbox(value="4BDU", label="PDB ID", placeholder="Enter PDB ID here...")
+    pdb_file = gr.File(label="Upload PDB/CIF File", visible=False)
+    visualize_btn = gr.Button("Visualize Structure")
+
+    molecule_output2 = Molecule3D(label="Protein Structure", reps=[
+        {
+            "model": 0,
+            "style": "cartoon",
+            "color": "whiteCarbon",
+            "residue_range": "",
+            "around": 0,
+            "byres": False,
+        }
+    ])
+
+    with gr.Row():
+        segment_input = gr.Textbox(value="A", label="Chain ID", placeholder="Enter Chain ID here...")
+        prediction_btn = gr.Button("Predict Binding Site")
+
+    molecule_output = gr.HTML(label="Protein Structure")
+    explanation_vis = gr.Markdown("""
+    Score dependent colorcoding:
+    - 0.0-0.2: white  
+    - 0.2–0.4: light orange  
+    - 0.4–0.6: orange
+    - 0.6–0.8: orangered
+    - 0.8–1.0: red
+    """)
+    predictions_output = gr.Textbox(label="Visualize Prediction with PyMol")
+    gr.Markdown("### Download:\n- List of predicted binding site residues\n- PDB with score in beta factor column")
+    download_output = gr.File(label="Download Files", file_count="multiple")
+    
+    def process_interface(mode, pdb_id, pdb_file, chain_id):
+        if mode == "PDB ID":
+            return process_pdb(pdb_id, chain_id)
+        elif mode == "Upload File":
+            _, ext = os.path.splitext(pdb_file.name)
+            file_path = os.path.join('./', f"{_}{ext}")
+            if ext == '.cif':
+                pdb_path = convert_cif_to_pdb(file_path)
+            else:
+                pdb_path= file_path
+            return process_pdb(pdb_path, chain_id)
+        else:
+            return "Error: Invalid mode selected", None, None
+
+    def fetch_interface(mode, pdb_id, pdb_file):
+        if mode == "PDB ID":
+            return fetch_pdb(pdb_id)
+        elif mode == "Upload File":
+            _, ext = os.path.splitext(pdb_file.name)
+            file_path = os.path.join('./', f"{_}{ext}")
+            #print(ext)
+            if ext == '.cif':
+                pdb_path = convert_cif_to_pdb(file_path)
+            else:
+                pdb_path= file_path
+            #print(pdb_path)
+            return pdb_path
+        else:
+            return "Error: Invalid mode selected"
+
+    def toggle_mode(selected_mode):
+        if selected_mode == "PDB ID":
+            return gr.update(visible=True), gr.update(visible=False)
+        else:
+            return gr.update(visible=False), gr.update(visible=True)
+
+    mode.change(
+        toggle_mode,
+        inputs=[mode],
+        outputs=[pdb_input, pdb_file]
+    )
+
+    prediction_btn.click(
+        process_interface, 
+        inputs=[mode, pdb_input, pdb_file, segment_input], 
+        outputs=[predictions_output, molecule_output, download_output]
+    )
+
+    visualize_btn.click(
+        fetch_interface, 
+        inputs=[mode, pdb_input, pdb_file], 
+        outputs=molecule_output2
+    )
+
+    gr.Markdown("## Examples")
+    gr.Examples(
+        examples=[
+            ["7RPZ", "A"],
+            ["2IWI", "B"],
+            ["2F6V", "A"]
+        ],
+        inputs=[pdb_input, segment_input],
+        outputs=[predictions_output, molecule_output, download_output]
+    )
+
+demo.launch(share=True)

test.ipynb

@@ -0,0 +1,1513 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "e776d9d6-417e-46d4-8061-846c055e1f8a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on local URL:  http://127.0.0.1:7873\n",
+      "* Running on public URL: https://120000a6aa9d78e04c.gradio.live\n",
+      "\n",
+      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"https://120000a6aa9d78e04c.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": []
+     },
+     "execution_count": 29,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from datetime import datetime\n",
+    "import gradio as gr\n",
+    "import requests\n",
+    "from Bio.PDB import PDBParser, MMCIFParser, PDBIO, Select\n",
+    "from Bio.PDB.Polypeptide import is_aa\n",
+    "from Bio.SeqUtils import seq1\n",
+    "from typing import Optional, Tuple\n",
+    "import numpy as np\n",
+    "import os\n",
+    "from gradio_molecule3d import Molecule3D\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\n",
+    "#from transformers import AutoTokenizer, DataCollatorForTokenClassification\n",
+    "\n",
+    "#from datasets import Dataset\n",
+    "\n",
+    "from scipy.special import expit\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_structure(pdb_id: str, output_dir: str = \".\") -> Optional[str]:\n",
+    "    \"\"\"\n",
+    "    Fetch the structure file for a given PDB ID. Prioritizes CIF files.\n",
+    "    If a structure file already exists locally, it uses that.\n",
+    "    \"\"\"\n",
+    "    file_path = download_structure(pdb_id, output_dir)\n",
+    "    if file_path:\n",
+    "        return file_path\n",
+    "    else:\n",
+    "        return None\n",
+    "\n",
+    "def download_structure(pdb_id: str, output_dir: str) -> Optional[str]:\n",
+    "    \"\"\"\n",
+    "    Attempt to download the structure file in CIF or PDB format.\n",
+    "    Returns the path to the downloaded file, or None if download fails.\n",
+    "    \"\"\"\n",
+    "    for ext in ['.cif', '.pdb']:\n",
+    "        file_path = os.path.join(output_dir, f\"{pdb_id}{ext}\")\n",
+    "        if os.path.exists(file_path):\n",
+    "            return file_path\n",
+    "        url = f\"https://files.rcsb.org/download/{pdb_id}{ext}\"\n",
+    "        try:\n",
+    "            response = requests.get(url, timeout=10)\n",
+    "            if response.status_code == 200:\n",
+    "                with open(file_path, 'wb') as f:\n",
+    "                    f.write(response.content)\n",
+    "                return file_path\n",
+    "        except Exception as e:\n",
+    "            print(f\"Download error for {pdb_id}{ext}: {e}\")\n",
+    "    return None\n",
+    "\n",
+    "def convert_cif_to_pdb(cif_path: str, output_dir: str = \".\") -> str:\n",
+    "    \"\"\"\n",
+    "    Convert a CIF file to PDB format using BioPython and return the PDB file path.\n",
+    "    \"\"\"\n",
+    "    pdb_path = os.path.join(output_dir, os.path.basename(cif_path).replace('.cif', '.pdb'))\n",
+    "    parser = MMCIFParser(QUIET=True)\n",
+    "    structure = parser.get_structure('protein', cif_path)\n",
+    "    io = PDBIO()\n",
+    "    io.set_structure(structure)\n",
+    "    io.save(pdb_path)\n",
+    "    return pdb_path\n",
+    "\n",
+    "def fetch_pdb(pdb_id):\n",
+    "    pdb_path = fetch_structure(pdb_id)\n",
+    "    if not pdb_path:\n",
+    "        return None\n",
+    "    _, ext = os.path.splitext(pdb_path)\n",
+    "    if ext == '.cif':\n",
+    "        pdb_path = convert_cif_to_pdb(pdb_path)\n",
+    "    return pdb_path\n",
+    "\n",
+    "def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list, protein_residues: list) -> str:\n",
+    "    \"\"\"\n",
+    "    Create a PDB file with only the selected chain and residues, replacing B-factor with prediction scores\n",
+    "    \"\"\"\n",
+    "    # Read the original PDB file\n",
+    "    parser = PDBParser(QUIET=True)\n",
+    "    structure = parser.get_structure('protein', input_pdb)\n",
+    "    \n",
+    "    # Prepare a new structure with only the specified chain and selected residues\n",
+    "    output_pdb = f\"{os.path.splitext(input_pdb)[0]}_{chain_id}_predictions_scores.pdb\"\n",
+    "    \n",
+    "    # Create scores dictionary for easy lookup\n",
+    "    scores_dict = {resi: score for resi, score in residue_scores}\n",
+    "\n",
+    "    # Create a custom Select class\n",
+    "    class ResidueSelector(Select):\n",
+    "        def __init__(self, chain_id, selected_residues, scores_dict):\n",
+    "            self.chain_id = chain_id\n",
+    "            self.selected_residues = selected_residues\n",
+    "            self.scores_dict = scores_dict\n",
+    "        \n",
+    "        def accept_chain(self, chain):\n",
+    "            return chain.id == self.chain_id\n",
+    "        \n",
+    "        def accept_residue(self, residue):\n",
+    "            return residue.id[1] in self.selected_residues\n",
+    "\n",
+    "        def accept_atom(self, atom):\n",
+    "            if atom.parent.id[1] in self.scores_dict:\n",
+    "                atom.bfactor = np.absolute(1-self.scores_dict[atom.parent.id[1]]) * 100\n",
+    "            return True\n",
+    "\n",
+    "    # Prepare output PDB with selected chain and residues, modified B-factors\n",
+    "    io = PDBIO()\n",
+    "    selector = ResidueSelector(chain_id, [res.id[1] for res in protein_residues], scores_dict)\n",
+    "    \n",
+    "    io.set_structure(structure[0])\n",
+    "    io.save(output_pdb, selector)\n",
+    "    \n",
+    "    return output_pdb\n",
+    "\n",
+    "def process_pdb(pdb_id_or_file, segment):\n",
+    "    # Determine if input is a PDB ID or file path\n",
+    "    if pdb_id_or_file.endswith('.pdb'):\n",
+    "        pdb_path = pdb_id_or_file\n",
+    "        pdb_id = os.path.splitext(os.path.basename(pdb_path))[0]\n",
+    "    else:\n",
+    "        pdb_id = pdb_id_or_file\n",
+    "        pdb_path = fetch_pdb(pdb_id)\n",
+    "    \n",
+    "    if not pdb_path:\n",
+    "        return \"Failed to fetch PDB file\", None, None\n",
+    "    \n",
+    "    # Determine the file format and choose the appropriate parser\n",
+    "    _, ext = os.path.splitext(pdb_path)\n",
+    "    parser = MMCIFParser(QUIET=True) if ext == '.cif' else PDBParser(QUIET=True)\n",
+    "    \n",
+    "    try:\n",
+    "        # Parse the structure file\n",
+    "        structure = parser.get_structure('protein', pdb_path)\n",
+    "    except Exception as e:\n",
+    "        return f\"Error parsing structure file: {e}\", None, None\n",
+    "    \n",
+    "    # Extract the specified chain\n",
+    "    try:\n",
+    "        chain = structure[0][segment]\n",
+    "    except KeyError:\n",
+    "        return \"Invalid Chain ID\", None, None\n",
+    "    \n",
+    "    protein_residues = [res for res in chain if is_aa(res)]\n",
+    "    sequence = \"\".join(seq1(res.resname) for res in protein_residues)\n",
+    "    sequence_id = [res.id[1] for res in protein_residues]\n",
+    "\n",
+    "    visualized_sequence = \"\".join(seq1(res.resname) for res in protein_residues)\n",
+    "    if sequence != visualized_sequence:\n",
+    "        raise ValueError(\"The visualized sequence does not match the prediction sequence\")\n",
+    "        \n",
+    "    scores = np.random.rand(len(sequence))\n",
+    "    normalized_scores = normalize_scores(scores)\n",
+    "    \n",
+    "    # Zip residues with scores to track the residue ID and score\n",
+    "    residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]\n",
+    "\n",
+    "    \n",
+    "    # Define the score brackets\n",
+    "    score_brackets = {\n",
+    "        \"0.0-0.2\": (0.0, 0.2),\n",
+    "        \"0.2-0.4\": (0.2, 0.4),\n",
+    "        \"0.4-0.6\": (0.4, 0.6),\n",
+    "        \"0.6-0.8\": (0.6, 0.8),\n",
+    "        \"0.8-1.0\": (0.8, 1.0)\n",
+    "    }\n",
+    "    \n",
+    "    # Initialize a dictionary to store residues by bracket\n",
+    "    residues_by_bracket = {bracket: [] for bracket in score_brackets}\n",
+    "    \n",
+    "    # Categorize residues into brackets\n",
+    "    for resi, score in residue_scores:\n",
+    "        for bracket, (lower, upper) in score_brackets.items():\n",
+    "            if lower <= score < upper:\n",
+    "                residues_by_bracket[bracket].append(resi)\n",
+    "                break\n",
+    "    \n",
+    "    # Preparing the result string\n",
+    "    current_time = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+    "    result_str = f\"Prediction for PDB: {pdb_id}, Chain: {segment}\\nDate: {current_time}\\n\\n\"\n",
+    "    result_str += \"Residues by Score Brackets:\\n\\n\"\n",
+    "    \n",
+    "    # Add residues for each bracket\n",
+    "    for bracket, residues in residues_by_bracket.items():\n",
+    "        result_str += f\"Bracket {bracket}:\\n\"\n",
+    "        result_str += \"Columns: Residue Name, Residue Number, One-letter Code, Normalized Score\\n\"\n",
+    "        result_str += \"\\n\".join([\n",
+    "            f\"{res.resname} {res.id[1]} {sequence[i]} {normalized_scores[i]:.2f}\" \n",
+    "            for i, res in enumerate(protein_residues) if res.id[1] in residues\n",
+    "        ])\n",
+    "        result_str += \"\\n\\n\"\n",
+    "\n",
+    "    # Create chain-specific PDB with scores in B-factor\n",
+    "    scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores, protein_residues)\n",
+    "\n",
+    "    # Molecule visualization with updated script with color mapping\n",
+    "    mol_vis = molecule(pdb_path, residue_scores, segment)#, color_map)\n",
+    "\n",
+    "    # Improved PyMOL command suggestions\n",
+    "    current_time = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+    "    pymol_commands = f\"Prediction for PDB: {pdb_id}, Chain: {segment}\\nDate: {current_time}\\n\\n\"\n",
+    "    \n",
+    "    pymol_commands += f\"\"\"\n",
+    "    # PyMOL Visualization Commands\n",
+    "    load {os.path.abspath(pdb_path)}, protein\n",
+    "    hide everything, all\n",
+    "    show cartoon, chain {segment}\n",
+    "    color white, chain {segment}\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    # Define colors for each score bracket\n",
+    "    bracket_colors = {\n",
+    "        \"0.0-0.2\": \"white\",\n",
+    "        \"0.2-0.4\": \"lightorange\",\n",
+    "        \"0.4-0.6\": \"orange\",\n",
+    "        \"0.6-0.8\": \"orangered\",\n",
+    "        \"0.8-1.0\": \"red\"\n",
+    "    }\n",
+    "    \n",
+    "    # Add PyMOL commands for each score bracket\n",
+    "    for bracket, residues in residues_by_bracket.items():\n",
+    "        if residues:  # Only add commands if there are residues in this bracket\n",
+    "            color = bracket_colors[bracket]\n",
+    "            resi_list = '+'.join(map(str, residues))\n",
+    "            pymol_commands += f\"\"\"\n",
+    "    select bracket_{bracket.replace('.', '').replace('-', '_')}, resi {resi_list} and chain {segment}\n",
+    "    show sticks, bracket_{bracket.replace('.', '').replace('-', '_')}\n",
+    "    color {color}, bracket_{bracket.replace('.', '').replace('-', '_')}\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    # Create prediction and scored PDB files\n",
+    "    prediction_file = f\"{pdb_id}_binding_site_residues.txt\"\n",
+    "    with open(prediction_file, \"w\") as f:\n",
+    "        f.write(result_str)\n",
+    "    \n",
+    "    return pymol_commands, mol_vis, [prediction_file,scored_pdb]\n",
+    "\n",
+    "def molecule(input_pdb, residue_scores=None, segment='A'):\n",
+    "    # More granular scoring for visualization\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",
+    "        # Filter residues based on their scores\n",
+    "        class1_score_residues = [resi for resi, score in residue_scores if 0.0 < score <= 0.2]\n",
+    "        class2_score_residues = [resi for resi, score in residue_scores if 0.2 < score <= 0.4]\n",
+    "        class3_score_residues = [resi for resi, score in residue_scores if 0.4 < score <= 0.6]\n",
+    "        class4_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.8]\n",
+    "        class5_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 1.0]\n",
+    "        \n",
+    "        high_score_script = \"\"\"\n",
+    "        // Load the original model and apply white cartoon style\n",
+    "        let chainModel = viewer.addModel(pdb, \"pdb\");\n",
+    "        chainModel.setStyle({}, {});\n",
+    "        chainModel.setStyle(\n",
+    "            {\"chain\": \"%s\"}, \n",
+    "            {\"cartoon\": {\"color\": \"white\"}}\n",
+    "        );\n",
+    "\n",
+    "        // Create a new model for high-scoring residues and apply red sticks style\n",
+    "        let class1Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        class1Model.setStyle({}, {});\n",
+    "        class1Model.setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
+    "            {\"stick\": {\"color\": \"0xFFFFFF\", \"opacity\": 0.5}}\n",
+    "        );\n",
+    "\n",
+    "        // Create a new model for high-scoring residues and apply red sticks style\n",
+    "        let class2Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        class2Model.setStyle({}, {});\n",
+    "        class2Model.setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
+    "            {\"stick\": {\"color\": \"0xFFD580\", \"opacity\": 0.7}}\n",
+    "        );\n",
+    "\n",
+    "        // Create a new model for high-scoring residues and apply red sticks style\n",
+    "        let class3Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        class3Model.setStyle({}, {});\n",
+    "        class3Model.setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
+    "            {\"stick\": {\"color\": \"0xFFA500\", \"opacity\": 1}}\n",
+    "        );\n",
+    "\n",
+    "        // Create a new model for high-scoring residues and apply red sticks style\n",
+    "        let class4Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        class4Model.setStyle({}, {});\n",
+    "        class4Model.setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
+    "            {\"stick\": {\"color\": \"0xFF4500\", \"opacity\": 1}}\n",
+    "        );\n",
+    "\n",
+    "        // Create a new model for high-scoring residues and apply red sticks style\n",
+    "        let class5Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        class5Model.setStyle({}, {});\n",
+    "        class5Model.setStyle(\n",
+    "            {\"chain\": \"%s\", \"resi\": [%s]}, \n",
+    "            {\"stick\": {\"color\": \"0xFF0000\", \"alpha\": 1}}\n",
+    "        );\n",
+    "\n",
+    "        \"\"\" % (\n",
+    "            segment,\n",
+    "            segment,\n",
+    "            \", \".join(str(resi) for resi in class1_score_residues),\n",
+    "            segment,\n",
+    "            \", \".join(str(resi) for resi in class2_score_residues),\n",
+    "            segment,\n",
+    "            \", \".join(str(resi) for resi in class3_score_residues),\n",
+    "            segment,\n",
+    "            \", \".join(str(resi) for resi in class4_score_residues),\n",
+    "            segment,\n",
+    "            \", \".join(str(resi) for resi in class5_score_residues)\n",
+    "        )\n",
+    "    \n",
+    "    # Generate the full HTML content\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",
+    "                \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: 18,\n",
+    "                                    padding: 4\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",
+    "# Gradio UI\n",
+    "with gr.Blocks(css=\"\"\"\n",
+    "    /* Customize Gradio button colors */\n",
+    "    #visualize-btn, #predict-btn {\n",
+    "        background-color: #FF7300; /* Deep orange */\n",
+    "        color: white;\n",
+    "        border-radius: 5px;\n",
+    "        padding: 10px;\n",
+    "        font-weight: bold;\n",
+    "    }\n",
+    "    #visualize-btn:hover, #predict-btn:hover {\n",
+    "        background-color: #CC5C00; /* Darkened orange on hover */\n",
+    "    }\n",
+    "\"\"\") as demo:\n",
+    "    gr.Markdown(\"# Protein Binding Site Prediction\")\n",
+    "    \n",
+    "    # Mode selection\n",
+    "    mode = gr.Radio(\n",
+    "        choices=[\"PDB ID\", \"Upload File\"],\n",
+    "        value=\"PDB ID\",\n",
+    "        label=\"Input Mode\",\n",
+    "        info=\"Choose whether to input a PDB ID or upload a PDB/CIF file.\"\n",
+    "    )\n",
+    "\n",
+    "    # Input components based on mode\n",
+    "    pdb_input = gr.Textbox(value=\"2F6V\", label=\"PDB ID\", placeholder=\"Enter PDB ID here...\")\n",
+    "    pdb_file = gr.File(label=\"Upload PDB/CIF File\", visible=False)\n",
+    "    visualize_btn = gr.Button(\"Visualize Structure\", elem_id=\"visualize-btn\")\n",
+    "\n",
+    "    molecule_output2 = Molecule3D(label=\"Protein Structure\", reps=[\n",
+    "        {\n",
+    "            \"model\": 0,\n",
+    "            \"style\": \"cartoon\",\n",
+    "            \"color\": \"whiteCarbon\",\n",
+    "            \"residue_range\": \"\",\n",
+    "            \"around\": 0,\n",
+    "            \"byres\": False,\n",
+    "        }\n",
+    "    ])\n",
+    "\n",
+    "    with gr.Row():\n",
+    "        segment_input = gr.Textbox(value=\"A\", label=\"Chain ID (protein)\", placeholder=\"Enter Chain ID here...\",\n",
+    "        info=\"Choose in which chain to predict binding sites.\")\n",
+    "        prediction_btn = gr.Button(\"Predict Binding Site\", elem_id=\"predict-btn\")\n",
+    "\n",
+    "    molecule_output = gr.HTML(label=\"Protein Structure\")\n",
+    "    explanation_vis = gr.Markdown(\"\"\"\n",
+    "    Score dependent colorcoding:\n",
+    "    - 0.0-0.2: white  \n",
+    "    - 0.2–0.4: light orange  \n",
+    "    - 0.4–0.6: orange\n",
+    "    - 0.6–0.8: orangered\n",
+    "    - 0.8–1.0: red\n",
+    "    \"\"\")\n",
+    "    predictions_output = gr.Textbox(label=\"Visualize Prediction with PyMol\")\n",
+    "    gr.Markdown(\"### Download:\\n- List of predicted binding site residues\\n- PDB with score in beta factor column\")\n",
+    "    download_output = gr.File(label=\"Download Files\", file_count=\"multiple\")\n",
+    "    \n",
+    "    def process_interface(mode, pdb_id, pdb_file, chain_id):\n",
+    "        if mode == \"PDB ID\":\n",
+    "            return process_pdb(pdb_id, chain_id)\n",
+    "        elif mode == \"Upload File\":\n",
+    "            _, ext = os.path.splitext(pdb_file.name)\n",
+    "            file_path = os.path.join('./', f\"{_}{ext}\")\n",
+    "            if ext == '.cif':\n",
+    "                pdb_path = convert_cif_to_pdb(file_path)\n",
+    "            else:\n",
+    "                pdb_path= file_path\n",
+    "            return process_pdb(pdb_path, chain_id)\n",
+    "        else:\n",
+    "            return \"Error: Invalid mode selected\", None, None\n",
+    "\n",
+    "    def fetch_interface(mode, pdb_id, pdb_file):\n",
+    "        if mode == \"PDB ID\":\n",
+    "            return fetch_pdb(pdb_id)\n",
+    "        elif mode == \"Upload File\":\n",
+    "            _, ext = os.path.splitext(pdb_file.name)\n",
+    "            file_path = os.path.join('./', f\"{_}{ext}\")\n",
+    "            #print(ext)\n",
+    "            if ext == '.cif':\n",
+    "                pdb_path = convert_cif_to_pdb(file_path)\n",
+    "            else:\n",
+    "                pdb_path= file_path\n",
+    "            #print(pdb_path)\n",
+    "            return pdb_path\n",
+    "        else:\n",
+    "            return \"Error: Invalid mode selected\"\n",
+    "\n",
+    "    def toggle_mode(selected_mode):\n",
+    "        if selected_mode == \"PDB ID\":\n",
+    "            return gr.update(visible=True), gr.update(visible=False)\n",
+    "        else:\n",
+    "            return gr.update(visible=False), gr.update(visible=True)\n",
+    "\n",
+    "    mode.change(\n",
+    "        toggle_mode,\n",
+    "        inputs=[mode],\n",
+    "        outputs=[pdb_input, pdb_file]\n",
+    "    )\n",
+    "\n",
+    "    prediction_btn.click(\n",
+    "        process_interface, \n",
+    "        inputs=[mode, pdb_input, pdb_file, segment_input], \n",
+    "        outputs=[predictions_output, molecule_output, download_output]\n",
+    "    )\n",
+    "\n",
+    "    visualize_btn.click(\n",
+    "        fetch_interface, \n",
+    "        inputs=[mode, pdb_input, pdb_file], \n",
+    "        outputs=molecule_output2\n",
+    "    )\n",
+    "\n",
+    "    gr.Markdown(\"## Examples\")\n",
+    "    gr.Examples(\n",
+    "        examples=[\n",
+    "            [\"7RPZ\", \"A\"],\n",
+    "            [\"2IWI\", \"B\"],\n",
+    "            [\"7LCJ\", \"R\"]\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": "440c87ed-45c9-4501-b208-409cbfd7858b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "d70c40b9-5d5a-4795-b2a2-149c4a57d16e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py:441: UserWarning: Examples will be cached but not all input components have example values. This may result in an exception being thrown by your function. If you do get an error while caching examples, make sure all of your inputs have example values for all of your examples or you provide default values for those particular parameters in your function.\n",
+      "  warnings.warn(\n",
+      "INFO:__main__:Using cached structure: ./7rpz.cif\n",
+      "INFO:__main__:Using cached structure: ./2iwi.cif\n",
+      "INFO:__main__:Using cached structure: ./2f6v.cif\n",
+      "INFO:httpx:HTTP Request: GET http://127.0.0.1:7862/gradio_api/startup-events \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on local URL:  http://127.0.0.1:7862\n",
+      "Caching examples at: '/home/frohlkin/Projects/LargeLanguageModels/Publication/test_webpage/.gradio/cached_examples/148'\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:httpx:HTTP Request: HEAD http://127.0.0.1:7862/ \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: GET https://api.gradio.app/pkg-version \"HTTP/1.1 200 OK\"\n",
+      "INFO:httpx:HTTP Request: GET https://api.gradio.app/v3/tunnel-request \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on public URL: https://de785d7cce806497e9.gradio.live\n",
+      "\n",
+      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:httpx:HTTP Request: HEAD https://de785d7cce806497e9.gradio.live \"HTTP/1.1 200 OK\"\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"https://de785d7cce806497e9.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Traceback (most recent call last):\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/routes.py\", line 990, in predict\n",
+      "    output = await route_utils.call_process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/route_utils.py\", line 322, in call_process_api\n",
+      "    output = await app.get_blocks().process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 2047, in process_api\n",
+      "    result = await self.call_function(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 1594, in call_function\n",
+      "    prediction = await anyio.to_thread.run_sync(  # type: ignore\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/to_thread.py\", line 56, in run_sync\n",
+      "    return await get_async_backend().run_sync_in_worker_thread(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 2405, in run_sync_in_worker_thread\n",
+      "    return await future\n",
+      "           ^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 914, in run\n",
+      "    result = context.run(func, *args)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/utils.py\", line 869, in wrapper\n",
+      "    response = f(*args, **kwargs)\n",
+      "               ^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 355, in load_example_with_output\n",
+      "    ) + self.load_from_cache(example_id)\n",
+      "        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 579, in load_from_cache\n",
+      "    output.append(component.read_from_flag(value_to_use))\n",
+      "                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/components/base.py\", line 366, in read_from_flag\n",
+      "    return self.data_model.from_json(json.loads(payload))\n",
+      "                                     ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/__init__.py\", line 346, in loads\n",
+      "    return _default_decoder.decode(s)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 337, in decode\n",
+      "    obj, end = self.raw_decode(s, idx=_w(s, 0).end())\n",
+      "               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 355, in raw_decode\n",
+      "    raise JSONDecodeError(\"Expecting value\", s, err.value) from None\n",
+      "json.decoder.JSONDecodeError: Expecting value: line 1 column 1 (char 0)\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/routes.py\", line 990, in predict\n",
+      "    output = await route_utils.call_process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/route_utils.py\", line 322, in call_process_api\n",
+      "    output = await app.get_blocks().process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 2047, in process_api\n",
+      "    result = await self.call_function(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 1594, in call_function\n",
+      "    prediction = await anyio.to_thread.run_sync(  # type: ignore\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/to_thread.py\", line 56, in run_sync\n",
+      "    return await get_async_backend().run_sync_in_worker_thread(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 2405, in run_sync_in_worker_thread\n",
+      "    return await future\n",
+      "           ^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 914, in run\n",
+      "    result = context.run(func, *args)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/utils.py\", line 869, in wrapper\n",
+      "    response = f(*args, **kwargs)\n",
+      "               ^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 355, in load_example_with_output\n",
+      "    ) + self.load_from_cache(example_id)\n",
+      "        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 579, in load_from_cache\n",
+      "    output.append(component.read_from_flag(value_to_use))\n",
+      "                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/components/base.py\", line 366, in read_from_flag\n",
+      "    return self.data_model.from_json(json.loads(payload))\n",
+      "                                     ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/__init__.py\", line 346, in loads\n",
+      "    return _default_decoder.decode(s)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 337, in decode\n",
+      "    obj, end = self.raw_decode(s, idx=_w(s, 0).end())\n",
+      "               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 355, in raw_decode\n",
+      "    raise JSONDecodeError(\"Expecting value\", s, err.value) from None\n",
+      "json.decoder.JSONDecodeError: Expecting value: line 1 column 1 (char 0)\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/routes.py\", line 990, in predict\n",
+      "    output = await route_utils.call_process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/route_utils.py\", line 322, in call_process_api\n",
+      "    output = await app.get_blocks().process_api(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 2047, in process_api\n",
+      "    result = await self.call_function(\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/blocks.py\", line 1594, in call_function\n",
+      "    prediction = await anyio.to_thread.run_sync(  # type: ignore\n",
+      "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/to_thread.py\", line 56, in run_sync\n",
+      "    return await get_async_backend().run_sync_in_worker_thread(\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 2405, in run_sync_in_worker_thread\n",
+      "    return await future\n",
+      "           ^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/anyio/_backends/_asyncio.py\", line 914, in run\n",
+      "    result = context.run(func, *args)\n",
+      "             ^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/utils.py\", line 869, in wrapper\n",
+      "    response = f(*args, **kwargs)\n",
+      "               ^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 355, in load_example_with_output\n",
+      "    ) + self.load_from_cache(example_id)\n",
+      "        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/helpers.py\", line 579, in load_from_cache\n",
+      "    output.append(component.read_from_flag(value_to_use))\n",
+      "                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/site-packages/gradio/components/base.py\", line 366, in read_from_flag\n",
+      "    return self.data_model.from_json(json.loads(payload))\n",
+      "                                     ^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/__init__.py\", line 346, in loads\n",
+      "    return _default_decoder.decode(s)\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 337, in decode\n",
+      "    obj, end = self.raw_decode(s, idx=_w(s, 0).end())\n",
+      "               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/home/frohlkin/anaconda3/envs/LLM/lib/python3.12/json/decoder.py\", line 355, in raw_decode\n",
+      "    raise JSONDecodeError(\"Expecting value\", s, err.value) from None\n",
+      "json.decoder.JSONDecodeError: Expecting value: line 1 column 1 (char 0)\n"
+     ]
+    }
+   ],
+   "source": [
+    "from datetime import datetime\n",
+    "import gradio as gr\n",
+    "import requests\n",
+    "from Bio.PDB import PDBParser, MMCIFParser, PDBIO, Select, Structure\n",
+    "from Bio.PDB.Polypeptide import is_aa\n",
+    "from Bio.SeqUtils import seq1\n",
+    "from typing import Optional, Tuple, Dict, List\n",
+    "import numpy as np\n",
+    "import os\n",
+    "from gradio_molecule3d import Molecule3D\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader\n",
+    "import re\n",
+    "import pandas as pd\n",
+    "import copy\n",
+    "from scipy.special import expit\n",
+    "import logging\n",
+    "import tempfile\n",
+    "\n",
+    "# Set up logging\n",
+    "logging.basicConfig(level=logging.INFO)\n",
+    "logger = logging.getLogger(__name__)\n",
+    "\n",
+    "class StructureError(Exception):\n",
+    "    \"\"\"Custom exception for structure-related errors\"\"\"\n",
+    "    pass\n",
+    "\n",
+    "def normalize_scores(scores: np.ndarray) -> np.ndarray:\n",
+    "    \"\"\"Normalize scores to range [0,1]\"\"\"\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: str) -> str:\n",
+    "    \"\"\"Read molecular structure file and return its content\"\"\"\n",
+    "    try:\n",
+    "        with open(pdb_path, 'r') as f:\n",
+    "            return f.read()\n",
+    "    except Exception as e:\n",
+    "        raise IOError(f\"Failed to read structure file: {e}\")\n",
+    "\n",
+    "def fetch_structure(pdb_id: str, output_dir: str = \".\") -> Optional[str]:\n",
+    "    \"\"\"Fetch structure file, trying multiple formats and sources\"\"\"\n",
+    "    try:\n",
+    "        # First try local cache\n",
+    "        for ext in ['.cif', '.pdb']:\n",
+    "            local_path = os.path.join(output_dir, f\"{pdb_id.lower()}{ext}\")\n",
+    "            if os.path.exists(local_path):\n",
+    "                logger.info(f\"Using cached structure: {local_path}\")\n",
+    "                return local_path\n",
+    "\n",
+    "        # Try different download sources\n",
+    "        sources = [\n",
+    "            f\"https://files.rcsb.org/download/{pdb_id.upper()}.cif\",\n",
+    "            f\"https://files.rcsb.org/download/{pdb_id.upper()}.pdb\",\n",
+    "            f\"https://files.rcsb.org/download/{pdb_id.lower()}.cif\",\n",
+    "            f\"https://files.rcsb.org/download/{pdb_id.lower()}.pdb\"\n",
+    "        ]\n",
+    "\n",
+    "        for url in sources:\n",
+    "            try:\n",
+    "                response = requests.get(url, timeout=10)\n",
+    "                if response.status_code == 200:\n",
+    "                    ext = '.cif' if 'cif' in url else '.pdb'\n",
+    "                    file_path = os.path.join(output_dir, f\"{pdb_id.lower()}{ext}\")\n",
+    "                    with open(file_path, 'wb') as f:\n",
+    "                        f.write(response.content)\n",
+    "                    logger.info(f\"Successfully downloaded: {url}\")\n",
+    "                    return file_path\n",
+    "            except Exception as e:\n",
+    "                logger.warning(f\"Failed to download from {url}: {e}\")\n",
+    "                continue\n",
+    "\n",
+    "        raise StructureError(f\"Failed to fetch structure for PDB ID: {pdb_id}\")\n",
+    "    except Exception as e:\n",
+    "        raise StructureError(f\"Error fetching structure: {e}\")\n",
+    "\n",
+    "def convert_cif_to_pdb(cif_path: str, output_dir: str = \".\") -> str:\n",
+    "    \"\"\"Convert CIF to PDB format with error handling\"\"\"\n",
+    "    try:\n",
+    "        pdb_path = os.path.join(output_dir, os.path.basename(cif_path).replace('.cif', '.pdb'))\n",
+    "        parser = MMCIFParser(QUIET=True)\n",
+    "        structure = parser.get_structure('protein', cif_path)\n",
+    "        io = PDBIO()\n",
+    "        io.set_structure(structure)\n",
+    "        io.save(pdb_path)\n",
+    "        return pdb_path\n",
+    "    except Exception as e:\n",
+    "        raise StructureError(f\"Failed to convert CIF to PDB: {e}\")\n",
+    "\n",
+    "def find_valid_chain(structure: Structure.Structure) -> Optional[str]:\n",
+    "    \"\"\"Find the first valid protein chain in the structure\"\"\"\n",
+    "    for model in structure:\n",
+    "        for chain in model:\n",
+    "            protein_residues = [res for res in chain if is_aa(res)]\n",
+    "            if len(protein_residues) > 0:\n",
+    "                return chain.id\n",
+    "    return None\n",
+    "\n",
+    "def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list, protein_residues: list) -> str:\n",
+    "    \"\"\"Create PDB file with selected chain and prediction scores in B-factor column\"\"\"\n",
+    "    class ResidueSelector(Select):\n",
+    "        def __init__(self, chain_id, selected_residues, scores_dict):\n",
+    "            self.chain_id = chain_id\n",
+    "            self.selected_residues = selected_residues\n",
+    "            self.scores_dict = scores_dict\n",
+    "        \n",
+    "        def accept_chain(self, chain):\n",
+    "            return chain.id == self.chain_id\n",
+    "        \n",
+    "        def accept_residue(self, residue):\n",
+    "            return residue.id[1] in self.selected_residues\n",
+    "\n",
+    "        def accept_atom(self, atom):\n",
+    "            if atom.parent.id[1] in self.scores_dict:\n",
+    "                atom.bfactor = np.absolute(1-self.scores_dict[atom.parent.id[1]]) * 100\n",
+    "            return True\n",
+    "\n",
+    "    try:\n",
+    "        parser = PDBParser(QUIET=True)\n",
+    "        structure = parser.get_structure('protein', input_pdb)\n",
+    "        output_pdb = f\"{os.path.splitext(input_pdb)[0]}_{chain_id}_predictions_scores.pdb\"\n",
+    "        scores_dict = {resi: score for resi, score in residue_scores}\n",
+    "        \n",
+    "        io = PDBIO()\n",
+    "        selector = ResidueSelector(chain_id, [res.id[1] for res in protein_residues], scores_dict)\n",
+    "        io.set_structure(structure[0])\n",
+    "        io.save(output_pdb, selector)\n",
+    "        \n",
+    "        return output_pdb\n",
+    "    except Exception as e:\n",
+    "        raise StructureError(f\"Failed to create chain-specific PDB: {e}\")\n",
+    "\n",
+    "def process_pdb(pdb_id_or_file: str, segment: str) -> Tuple[str, str, List[str]]:\n",
+    "    \"\"\"Process PDB/CIF file and generate visualizations and predictions\"\"\"\n",
+    "    try:\n",
+    "        # Handle input\n",
+    "        if pdb_id_or_file.endswith(('.pdb', '.cif')):\n",
+    "            pdb_path = pdb_id_or_file\n",
+    "            pdb_id = os.path.splitext(os.path.basename(pdb_path))[0]\n",
+    "        else:\n",
+    "            pdb_id = pdb_id_or_file\n",
+    "            pdb_path = fetch_structure(pdb_id)\n",
+    "\n",
+    "        if not pdb_path:\n",
+    "            raise StructureError(\"Failed to obtain structure file\")\n",
+    "\n",
+    "        # Parse structure\n",
+    "        parser = MMCIFParser(QUIET=True) if pdb_path.endswith('.cif') else PDBParser(QUIET=True)\n",
+    "        structure = parser.get_structure('protein', pdb_path)\n",
+    "\n",
+    "        # Handle chain selection\n",
+    "        if segment == 'auto' or not segment:\n",
+    "            segment = find_valid_chain(structure)\n",
+    "            if not segment:\n",
+    "                raise StructureError(\"No valid protein chains found in structure\")\n",
+    "        \n",
+    "        try:\n",
+    "            chain = structure[0][segment]\n",
+    "        except KeyError:\n",
+    "            valid_chain = find_valid_chain(structure)\n",
+    "            if valid_chain:\n",
+    "                chain = structure[0][valid_chain]\n",
+    "                segment = valid_chain\n",
+    "                logger.info(f\"Using alternative chain {segment}\")\n",
+    "            else:\n",
+    "                raise StructureError(f\"Invalid chain ID '{segment}'. Structure has no valid protein chains.\")\n",
+    "\n",
+    "        # Process chain\n",
+    "        protein_residues = [res for res in chain if is_aa(res)]\n",
+    "        if not protein_residues:\n",
+    "            raise StructureError(f\"No amino acid residues found in chain {segment}\")\n",
+    "\n",
+    "        sequence = \"\".join(seq1(res.resname) for res in protein_residues)\n",
+    "        sequence_id = [res.id[1] for res in protein_residues]\n",
+    "        \n",
+    "        # Generate predictions (currently random)\n",
+    "        scores = np.random.rand(len(sequence))\n",
+    "        normalized_scores = normalize_scores(scores)\n",
+    "        residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]\n",
+    "\n",
+    "        # Generate outputs\n",
+    "        result_str = generate_results_string(pdb_id, segment, protein_residues, normalized_scores, sequence)\n",
+    "        scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores, protein_residues)\n",
+    "        mol_vis = molecule(pdb_path, residue_scores, segment)\n",
+    "        pymol_commands = generate_pymol_commands(pdb_id, segment, residue_scores, pdb_path)\n",
+    "\n",
+    "        # Save results\n",
+    "        prediction_file = f\"{pdb_id}_binding_site_residues.txt\"\n",
+    "        with open(prediction_file, \"w\") as f:\n",
+    "            f.write(result_str)\n",
+    "\n",
+    "        return pymol_commands, mol_vis, [prediction_file, scored_pdb]\n",
+    "\n",
+    "    except StructureError as e:\n",
+    "        return str(e), None, None\n",
+    "    except Exception as e:\n",
+    "        return f\"An unexpected error occurred: {str(e)}\", None, None\n",
+    "\n",
+    "def generate_results_string(pdb_id: str, segment: str, protein_residues: list, \n",
+    "                          normalized_scores: np.ndarray, sequence: str) -> str:\n",
+    "    \"\"\"Generate formatted results string with predictions\"\"\"\n",
+    "    score_brackets = {\n",
+    "        \"0.0-0.2\": (0.0, 0.2),\n",
+    "        \"0.2-0.4\": (0.2, 0.4),\n",
+    "        \"0.4-0.6\": (0.4, 0.6),\n",
+    "        \"0.6-0.8\": (0.6, 0.8),\n",
+    "        \"0.8-1.0\": (0.8, 1.0)\n",
+    "    }\n",
+    "    \n",
+    "    residues_by_bracket = {bracket: [] for bracket in score_brackets}\n",
+    "    \n",
+    "    # Categorize residues\n",
+    "    for i, score in enumerate(normalized_scores):\n",
+    "        for bracket, (lower, upper) in score_brackets.items():\n",
+    "            if lower <= score < upper:\n",
+    "                residues_by_bracket[bracket].append(protein_residues[i])\n",
+    "                break\n",
+    "    \n",
+    "    # Format results\n",
+    "    current_time = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+    "    result_str = f\"\"\"Prediction Results\n",
+    "========================\n",
+    "PDB: {pdb_id}\n",
+    "Chain: {segment}\n",
+    "Date: {current_time}\n",
+    "\n",
+    "Analysis by Score Brackets\n",
+    "========================\n",
+    "\"\"\"\n",
+    "    \n",
+    "    for bracket, residues in residues_by_bracket.items():\n",
+    "        if residues:  # Only show brackets with residues\n",
+    "            result_str += f\"\\nBracket {bracket}:\\n\"\n",
+    "            result_str += \"ResName ResNum Code Score\\n\"\n",
+    "            result_str += \"-\" * 30 + \"\\n\"\n",
+    "            result_str += \"\\n\".join([\n",
+    "                f\"{res.resname:6} {res.id[1]:6} {sequence[i]:4} {normalized_scores[i]:6.2f}\" \n",
+    "                for i, res in enumerate(protein_residues) if res in residues\n",
+    "            ])\n",
+    "            result_str += \"\\n\"\n",
+    "    \n",
+    "    return result_str\n",
+    "\n",
+    "def generate_pymol_commands(pdb_id: str, segment: str, residue_scores: list, pdb_path: str) -> str:\n",
+    "    \"\"\"Generate PyMOL visualization commands\"\"\"\n",
+    "    # Group residues by score ranges\n",
+    "    score_groups = {\n",
+    "        \"very_low\": [], \"low\": [], \"medium\": [], \"high\": [], \"very_high\": []\n",
+    "    }\n",
+    "    \n",
+    "    for resi, score in residue_scores:\n",
+    "        if score <= 0.2:\n",
+    "            score_groups[\"very_low\"].append(str(resi))\n",
+    "        elif score <= 0.4:\n",
+    "            score_groups[\"low\"].append(str(resi))\n",
+    "        elif score <= 0.6:\n",
+    "            score_groups[\"medium\"].append(str(resi))\n",
+    "        elif score <= 0.8:\n",
+    "            score_groups[\"high\"].append(str(resi))\n",
+    "        else:\n",
+    "            score_groups[\"very_high\"].append(str(resi))\n",
+    "\n",
+    "    current_time = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+    "    commands = f\"\"\"# PyMOL Script for {pdb_id} Chain {segment}\n",
+    "# Generated: {current_time}\n",
+    "\n",
+    "# Load structure and set initial display\n",
+    "load {os.path.abspath(pdb_path)}, protein\n",
+    "bg_color white\n",
+    "hide everything\n",
+    "show cartoon, chain {segment}\n",
+    "color white, chain {segment}\n",
+    "\n",
+    "# Create selection groups by score\n",
+    "\"\"\"\n",
+    "    \n",
+    "    color_scheme = {\n",
+    "        \"very_low\": \"white\",\n",
+    "        \"low\": \"lightorange\",\n",
+    "        \"medium\": \"orange\",\n",
+    "        \"high\": \"orangered\",\n",
+    "        \"very_high\": \"red\"\n",
+    "    }\n",
+    "    \n",
+    "    for group, residues in score_groups.items():\n",
+    "        if residues:\n",
+    "            resi_str = \"+\".join(residues)\n",
+    "            commands += f\"\"\"\n",
+    "# {group.replace('_', ' ').title()} scoring residues\n",
+    "select {group}, chain {segment} and resi {resi_str}\n",
+    "show sticks, {group}\n",
+    "color {color_scheme[group]}, {group}\"\"\"\n",
+    "    \n",
+    "    commands += \"\"\"\n",
+    "\n",
+    "# Center and zoom\n",
+    "center chain {}\n",
+    "zoom chain {}\n",
+    "\"\"\"\n",
+    "\n",
+    "    return commands\n",
+    "\n",
+    "def molecule(input_pdb: str, residue_scores: list = None, segment: str = 'A') -> str:\n",
+    "    \"\"\"Generate interactive 3D molecule visualization\"\"\"\n",
+    "    try:\n",
+    "        mol = read_mol(input_pdb)\n",
+    "    except Exception as e:\n",
+    "        return f'<div class=\"error\">Error loading structure: {str(e)}</div>'\n",
+    "\n",
+    "    # Prepare residue groups for visualization\n",
+    "    vis_groups = {\n",
+    "        \"class1\": [],  # 0.0-0.2\n",
+    "        \"class2\": [],  # 0.2-0.4\n",
+    "        \"class3\": [],  # 0.4-0.6\n",
+    "        \"class4\": [],  # 0.6-0.8\n",
+    "        \"class5\": []   # 0.8-1.0\n",
+    "    }\n",
+    "\n",
+    "    if residue_scores:\n",
+    "        for resi, score in residue_scores:\n",
+    "            if score <= 0.2:\n",
+    "                vis_groups[\"class1\"].append(resi)\n",
+    "            elif score <= 0.4:\n",
+    "                vis_groups[\"class2\"].append(resi)\n",
+    "            elif score <= 0.6:\n",
+    "                vis_groups[\"class3\"].append(resi)\n",
+    "            elif score <= 0.8:\n",
+    "                vis_groups[\"class4\"].append(resi)\n",
+    "            else:\n",
+    "                vis_groups[\"class5\"].append(resi)\n",
+    "\n",
+    "    # Generate visualization script\n",
+    "    vis_script = f\"\"\"\n",
+    "        // Base model setup\n",
+    "        let chainModel = viewer.addModel(pdb, \"pdb\");\n",
+    "        chainModel.setStyle({{}}, {{}});\n",
+    "        chainModel.setStyle(\n",
+    "            {{\"chain\": \"{segment}\"}}, \n",
+    "            {{\"cartoon\": {{\"color\": \"white\"}}}}\n",
+    "        );\n",
+    "    \"\"\"\n",
+    "\n",
+    "    # Color schemes for different score ranges\n",
+    "    color_schemes = {\n",
+    "        \"class1\": {\"color\": \"0xFFFFFF\", \"opacity\": 0.5},  # White\n",
+    "        \"class2\": {\"color\": \"0xFFD580\", \"opacity\": 0.7},  # Light orange\n",
+    "        \"class3\": {\"color\": \"0xFFA500\", \"opacity\": 1.0},  # Orange\n",
+    "        \"class4\": {\"color\": \"0xFF4500\", \"opacity\": 1.0},  # Orange red\n",
+    "        \"class5\": {\"color\": \"0xFF0000\", \"opacity\": 1.0}   # Red\n",
+    "    }\n",
+    "\n",
+    "    # Add visualization for each group\n",
+    "    for group, residues in vis_groups.items():\n",
+    "        if residues:\n",
+    "            color_scheme = color_schemes[group]\n",
+    "            vis_script += f\"\"\"\n",
+    "        let {group}Model = viewer.addModel(pdb, \"pdb\");\n",
+    "        {group}Model.setStyle({{}}, {{}});\n",
+    "        {group}Model.setStyle(\n",
+    "            {{\"chain\": \"{segment}\", \"resi\": [{\", \".join(map(str, residues))}]}},\n",
+    "            {{\"stick\": {{\"color\": \"{color_scheme[\"color\"]}\", \"opacity\": {color_scheme[\"opacity\"]}}}}}\n",
+    "        );\n",
+    "        \"\"\"\n",
+    "\n",
+    "    # Generate full HTML with enhanced controls and information\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",
+    "        .controls {{\n",
+    "            position: absolute;\n",
+    "            top: 10px;\n",
+    "            left: 10px;\n",
+    "            background: rgba(255, 255, 255, 0.8);\n",
+    "            padding: 10px;\n",
+    "            border-radius: 5px;\n",
+    "            z-index: 1000;\n",
+    "        }}\n",
+    "        .legend {{\n",
+    "            position: absolute;\n",
+    "            bottom: 10px;\n",
+    "            right: 10px;\n",
+    "            background: rgba(255, 255, 255, 0.8);\n",
+    "            padding: 10px;\n",
+    "            border-radius: 5px;\n",
+    "            z-index: 1000;\n",
+    "        }}\n",
+    "        .error {{\n",
+    "            color: red;\n",
+    "            padding: 20px;\n",
+    "            text-align: center;\n",
+    "            font-weight: bold;\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\">\n",
+    "            <div class=\"controls\">\n",
+    "                <button onclick=\"toggleStyle('cartoon')\">Toggle Cartoon</button>\n",
+    "                <button onclick=\"toggleStyle('stick')\">Toggle Sticks</button>\n",
+    "                <button onclick=\"resetView()\">Reset View</button>\n",
+    "                <button onclick=\"toggleSpin()\">Toggle Spin</button>\n",
+    "            </div>\n",
+    "            <div class=\"legend\">\n",
+    "                <div><span style=\"color: #FF0000\">■</span> Very High (0.8-1.0)</div>\n",
+    "                <div><span style=\"color: #FF4500\">■</span> High (0.6-0.8)</div>\n",
+    "                <div><span style=\"color: #FFA500\">■</span> Medium (0.4-0.6)</div>\n",
+    "                <div><span style=\"color: #FFD580\">■</span> Low (0.2-0.4)</div>\n",
+    "                <div><span style=\"color: #FFFFFF\">■</span> Very Low (0.0-0.2)</div>\n",
+    "            </div>\n",
+    "        </div>\n",
+    "        <script>\n",
+    "            let pdb = `{mol}`;\n",
+    "            let viewer;\n",
+    "            let isSpinning = false;\n",
+    "\n",
+    "            $(document).ready(function () {{\n",
+    "                let element = $(\"#container\");\n",
+    "                let config = {{ backgroundColor: \"white\" }};\n",
+    "                viewer = $3Dmol.createViewer(element, config);\n",
+    "                \n",
+    "                {vis_script}\n",
+    "                \n",
+    "                // Enhanced 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: 18,\n",
+    "                                    padding: 4\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",
+    "\n",
+    "            function toggleStyle(style) {{\n",
+    "                let elements = viewer.selectedAtoms({{}});\n",
+    "                let currentStyle = elements.style[style];\n",
+    "                elements.style[style] = !currentStyle;\n",
+    "                viewer.render();\n",
+    "            }}\n",
+    "\n",
+    "            function resetView() {{\n",
+    "                viewer.zoomTo();\n",
+    "                viewer.render();\n",
+    "            }}\n",
+    "\n",
+    "            function toggleSpin() {{\n",
+    "                isSpinning = !isSpinning;\n",
+    "                viewer.spin(isSpinning);\n",
+    "            }}\n",
+    "        </script>\n",
+    "    </body>\n",
+    "    </html>\n",
+    "    \"\"\"\n",
+    "    \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",
+    "def create_ui():\n",
+    "    with gr.Blocks(title=\"Protein Binding Site Prediction\", theme=gr.themes.Base()) as demo:\n",
+    "        gr.Markdown(\"\"\"\n",
+    "        # Protein Binding Site Prediction\n",
+    "        \n",
+    "        This tool helps you visualize and analyze potential binding sites in protein structures.\n",
+    "        You can either:\n",
+    "        1. Enter a PDB ID (e.g., \"4BDU\")\n",
+    "        2. Upload your own PDB/CIF file\n",
+    "        \n",
+    "        The tool will analyze the structure and show predictions using a color gradient from white (low probability) to red (high probability).\n",
+    "        \"\"\")\n",
+    "        \n",
+    "        with gr.Row():\n",
+    "            with gr.Column(scale=2):\n",
+    "                # Input components\n",
+    "                mode = gr.Radio(\n",
+    "                    choices=[\"PDB ID\", \"Upload File\"],\n",
+    "                    value=\"PDB ID\",\n",
+    "                    label=\"Input Mode\",\n",
+    "                    info=\"Choose whether to input a PDB ID or upload a PDB/CIF file\"\n",
+    "                )\n",
+    "                \n",
+    "                with gr.Group():\n",
+    "                    pdb_input = gr.Textbox(\n",
+    "                        value=\"4BDU\",\n",
+    "                        label=\"PDB ID\",\n",
+    "                        placeholder=\"Enter PDB ID (e.g., 4BDU)\",\n",
+    "                        info=\"Enter a valid PDB ID from the Protein Data Bank\"\n",
+    "                    )\n",
+    "                    pdb_file = gr.File(\n",
+    "                        label=\"Upload PDB/CIF File\",\n",
+    "                        file_types=[\".pdb\", \".cif\"],\n",
+    "                        visible=False\n",
+    "                    )\n",
+    "                \n",
+    "                segment_input = gr.Textbox(\n",
+    "                    value=\"A\",\n",
+    "                    label=\"Chain ID\",\n",
+    "                    placeholder=\"Enter Chain ID or leave empty for automatic selection\",\n",
+    "                    info=\"Specify which protein chain to analyze, or leave empty for automatic selection\"\n",
+    "                )\n",
+    "\n",
+    "            with gr.Column(scale=1):\n",
+    "                visualize_btn = gr.Button(\"Visualize Structure\", variant=\"primary\")\n",
+    "                prediction_btn = gr.Button(\"Predict Binding Site\", variant=\"secondary\")\n",
+    "                \n",
+    "                gr.Markdown(\"\"\"\n",
+    "                ### Color Legend\n",
+    "                - White: Very Low (0.0-0.2)\n",
+    "                - Light Orange: Low (0.2-0.4)\n",
+    "                - Orange: Medium (0.4-0.6)\n",
+    "                - Orange Red: High (0.6-0.8)\n",
+    "                - Red: Very High (0.8-1.0)\n",
+    "                \"\"\")\n",
+    "\n",
+    "        with gr.Tab(\"3D Visualization\"):\n",
+    "            molecule_output = gr.HTML(label=\"Interactive 3D Structure\")\n",
+    "            \n",
+    "        with gr.Tab(\"Analysis Results\"):\n",
+    "            predictions_output = gr.Textbox(\n",
+    "                label=\"PyMOL Visualization Commands\",\n",
+    "                info=\"Copy these commands into PyMOL to recreate the visualization\"\n",
+    "            )\n",
+    "            download_output = gr.File(\n",
+    "                label=\"Download Results\",\n",
+    "                file_count=\"multiple\"\n",
+    "            )\n",
+    "\n",
+    "        # Error message container\n",
+    "        error_output = gr.Markdown(visible=False)\n",
+    "\n",
+    "        # Mode change handler\n",
+    "        def toggle_mode(selected_mode):\n",
+    "            return {\n",
+    "                pdb_input: gr.update(visible=selected_mode == \"PDB ID\"),\n",
+    "                pdb_file: gr.update(visible=selected_mode == \"Upload File\")\n",
+    "            }\n",
+    "\n",
+    "        mode.change(\n",
+    "            toggle_mode,\n",
+    "            inputs=[mode],\n",
+    "            outputs=[pdb_input, pdb_file]\n",
+    "        )\n",
+    "\n",
+    "        # Process handlers\n",
+    "        def handle_visualization(mode, pdb_id, pdb_file):\n",
+    "            try:\n",
+    "                result = fetch_interface(mode, pdb_id, pdb_file)\n",
+    "                if isinstance(result, str) and result.startswith(\"Error\"):\n",
+    "                    return None, gr.update(visible=True, value=f\"```\\n{result}\\n```\")\n",
+    "                return result, gr.update(visible=False)\n",
+    "            except Exception as e:\n",
+    "                return None, gr.update(visible=True, value=f\"```\\nError: {str(e)}\\n```\")\n",
+    "\n",
+    "        def handle_prediction(mode, pdb_id, pdb_file, chain_id):\n",
+    "            try:\n",
+    "                predictions, vis, downloads = process_interface(mode, pdb_id, pdb_file, chain_id)\n",
+    "                if isinstance(predictions, str) and \"Error\" in predictions:\n",
+    "                    return (\n",
+    "                        None,\n",
+    "                        None,\n",
+    "                        None,\n",
+    "                        gr.update(visible=True, value=f\"```\\n{predictions}\\n```\")\n",
+    "                    )\n",
+    "                return (\n",
+    "                    predictions,\n",
+    "                    vis,\n",
+    "                    downloads,\n",
+    "                    gr.update(visible=False)\n",
+    "                )\n",
+    "            except Exception as e:\n",
+    "                error_msg = f\"\"\"Error processing structure:\n",
+    "```\n",
+    "{str(e)}\n",
+    "\n",
+    "Troubleshooting tips:\n",
+    "1. Check if the PDB ID is valid\n",
+    "2. Ensure the Chain ID exists in the structure\n",
+    "3. Try leaving Chain ID empty for automatic selection\n",
+    "4. If uploading a file, ensure it's a valid PDB/CIF format\n",
+    "```\"\"\"\n",
+    "                return None, None, None, gr.update(visible=True, value=error_msg)\n",
+    "\n",
+    "        def fetch_interface(mode, pdb_id, pdb_file):\n",
+    "            try:\n",
+    "                if mode == \"PDB ID\":\n",
+    "                    if not pdb_id or len(pdb_id.strip()) != 4:\n",
+    "                        raise ValueError(\"Please enter a valid 4-character PDB ID\")\n",
+    "                    return fetch_pdb(pdb_id.strip())\n",
+    "                elif mode == \"Upload File\":\n",
+    "                    if not pdb_file:\n",
+    "                        raise ValueError(\"Please upload a PDB or CIF file\")\n",
+    "                    _, ext = os.path.splitext(pdb_file.name)\n",
+    "                    if ext.lower() not in ['.pdb', '.cif']:\n",
+    "                        raise ValueError(\"Only .pdb and .cif files are supported\")\n",
+    "                    \n",
+    "                    # Create temp directory for file handling\n",
+    "                    with tempfile.TemporaryDirectory() as temp_dir:\n",
+    "                        temp_path = os.path.join(temp_dir, os.path.basename(pdb_file.name))\n",
+    "                        with open(temp_path, 'wb') as f:\n",
+    "                            f.write(pdb_file.read())\n",
+    "                        \n",
+    "                        if ext.lower() == '.cif':\n",
+    "                            return convert_cif_to_pdb(temp_path)\n",
+    "                        return temp_path\n",
+    "                else:\n",
+    "                    raise ValueError(\"Invalid mode selected\")\n",
+    "            except Exception as e:\n",
+    "                return f\"Error: {str(e)}\"\n",
+    "\n",
+    "        # Connect event handlers\n",
+    "        visualize_btn.click(\n",
+    "            handle_visualization,\n",
+    "            inputs=[mode, pdb_input, pdb_file],\n",
+    "            outputs=[molecule_output, error_output]\n",
+    "        )\n",
+    "\n",
+    "        prediction_btn.click(\n",
+    "            handle_prediction,\n",
+    "            inputs=[mode, pdb_input, pdb_file, segment_input],\n",
+    "            outputs=[predictions_output, molecule_output, download_output, error_output]\n",
+    "        )\n",
+    "\n",
+    "        # Add examples\n",
+    "        gr.Examples(\n",
+    "            examples=[\n",
+    "                [\"PDB ID\", \"7RPZ\", None, \"A\"],\n",
+    "                [\"PDB ID\", \"2IWI\", None, \"B\"],\n",
+    "                [\"PDB ID\", \"2F6V\", None, \"A\"]\n",
+    "            ],\n",
+    "            inputs=[mode, pdb_input, pdb_file, segment_input],\n",
+    "            outputs=[predictions_output, molecule_output, download_output, error_output],\n",
+    "            fn=handle_prediction,\n",
+    "            cache_examples=True\n",
+    "        )\n",
+    "\n",
+    "        # Add documentation\n",
+    "        gr.Markdown(\"\"\"\n",
+    "        ## Usage Instructions\n",
+    "        \n",
+    "        1. **Input Structure:**\n",
+    "           - Enter a PDB ID (e.g., \"4BDU\") or upload your own structure file\n",
+    "           - The tool supports both PDB (.pdb) and mmCIF (.cif) formats\n",
+    "        \n",
+    "        2. **Select Chain:**\n",
+    "           - Enter a specific chain ID (e.g., \"A\")\n",
+    "           - Leave empty for automatic selection of the first valid protein chain\n",
+    "        \n",
+    "        3. **Analyze:**\n",
+    "           - Click \"Visualize Structure\" to view the 3D structure\n",
+    "           - Click \"Predict Binding Site\" to perform binding site analysis\n",
+    "        \n",
+    "        4. **Results:**\n",
+    "           - Interactive 3D visualization with color-coded predictions\n",
+    "           - PyMOL commands for external visualization\n",
+    "           - Downloadable results files\n",
+    "        \n",
+    "        ## Troubleshooting\n",
+    "        \n",
+    "        If you encounter issues:\n",
+    "        1. Ensure your PDB ID is valid and exists in the PDB database\n",
+    "        2. Check that your uploaded file is a valid PDB/CIF format\n",
+    "        3. Try automatic chain selection if your specified chain isn't found\n",
+    "        4. Clear your browser cache if visualizations don't load\n",
+    "        \"\"\")\n",
+    "\n",
+    "    return demo\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    demo = create_ui()\n",
+    "    demo.launch(share=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9125d1c4-e2ae-4e40-ba36-7ae944512b8e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "85c0728a-a15b-4118-b920-5f55a2f5f79a",
+   "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.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}