Update

.ipynb_checkpoints/app-checkpoint.py

@@ -37,7 +37,7 @@ from scipy.special import expit
 
 import requests
 
-import py3Dmol
+from gradio_molecule3d import Molecule3D
 
 #import peft
 #from peft import get_peft_config, PeftModel, PeftConfig, inject_adapter_in_model, LoraConfig
@@ -46,6 +46,21 @@ import py3Dmol
 checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
 max_length = 1500
 
+# Default representations for molecule rendering
+reps = [
+    {
+        "model": 0,
+        "chain": "",
+        "resname": "",
+        "style": "cartoon",
+        "color": "spectrum",
+        "residue_range": "",
+        "around": 0,
+        "byres": False,
+        "visible": True
+    }
+]
+
 # Load model and move to device
 model, tokenizer = load_model(checkpoint, max_length)
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -53,7 +68,6 @@ model.to(device)
 model.eval()
 
 def create_dataset(tokenizer, seqs, labels, checkpoint):
-    
     tokenized = tokenizer(seqs, max_length=max_length, padding=False, truncation=True)
     dataset = Dataset.from_dict(tokenized)
     
@@ -68,7 +82,6 @@ def create_dataset(tokenizer, seqs, labels, checkpoint):
     return dataset
 
 def convert_predictions(input_logits):
-    
     all_probs = []
     for logits in input_logits:
         logits = logits.reshape(-1, 2)
@@ -78,13 +91,11 @@ def convert_predictions(input_logits):
     return np.concatenate(all_probs)
 
 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 predict_protein_sequence(test_one_letter_sequence):    
-    
     # Sanitize input sequence
     test_one_letter_sequence = test_one_letter_sequence.replace("O", "X") \
         .replace("B", "X").replace("U", "X") \
@@ -135,58 +146,88 @@ def predict_protein_sequence(test_one_letter_sequence):
     
     return result_str
 
-def fetch_and_display_pdb(pdb_id):
-    
+def fetch_pdb(pdb_id):
     try:
+        # Create a directory to store PDB files if it doesn't exist
+        os.makedirs('pdb_files', exist_ok=True)
+        
         # Fetch the PDB structure from RCSB
         pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'
-        response = requests.get(pdb_url)
-        
-        if response.status_code != 200:
-            return "Failed to load PDB structure. Please check the PDB ID."
+        pdb_path = f'pdb_files/{pdb_id}.pdb'
         
-        pdb_structure = response.text
+        # Download the file
+        response = requests.get(pdb_url)
         
-        # Prepare the 3D molecular visualization
-        visualization = f"""
-        <div id="container" style="width: 100%; height: 400px; position: relative;"></div>
-        <script src="https://3dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
-        <script>
-            let viewer = $3Dmol.createViewer(document.getElementById("container"));
-            viewer.addModel(`{pdb_structure}`, "pdb");
-            viewer.setStyle({{}}, {{"cartoon": {{"color": "spectrum"}}}});
-            viewer.zoomTo();
-            viewer.render();
-        </script>
-        """
-        return visualization
+        if response.status_code == 200:
+            with open(pdb_path, 'wb') as f:
+                f.write(response.content)
+            return pdb_path
+        else:
+            return None
     
     except Exception as e:
-        return f"Error visualizing PDB: {str(e)}"
+        print(f"Error fetching PDB: {e}")
+        return None
 
-def gradio_interface(sequence, pdb_id):
-    
+def process_input(sequence, pdb_id):
     # Predict binding sites
     binding_site_predictions = predict_protein_sequence(sequence)
     
-    # Fetch and visualize PDB structure
-    pdb_structure_html = fetch_and_display_pdb(pdb_id)
+    # Fetch PDB file
+    pdb_path = fetch_pdb(pdb_id)
     
-    return binding_site_predictions, pdb_structure_html
+    return binding_site_predictions, pdb_path
 
 # Create Gradio interface
-interface = gr.Interface(
-    fn=gradio_interface,
-    inputs=[
-        gr.Textbox(lines=2, placeholder="Enter protein sequence here...", label="Protein Sequence"),
-        gr.Textbox(lines=1, placeholder="Enter PDB ID here...", label="PDB ID for 3D Visualization")
-    ],
-    outputs=[
-        gr.Textbox(label="Binding Site Predictions"),
-        gr.HTML(label="3D Molecular Viewer")
-    ],
-    title="Protein Binding Site Prediction and 3D Structure Viewer",
-    description="Input a protein sequence to predict binding sites and view the protein structure in 3D using its PDB ID.",
-)
-
-interface.launch()
+with gr.Blocks() as demo:
+    gr.Markdown("# Protein Binding Site Prediction")
+    
+    with gr.Row():
+        with gr.Column():
+            # Sequence input
+            sequence_input = gr.Textbox(
+                lines=2, 
+                placeholder="Enter protein sequence here...", 
+                label="Protein Sequence"
+            )
+            
+            # PDB ID input
+            pdb_input = gr.Textbox(
+                lines=1, 
+                placeholder="Enter PDB ID here...", 
+                label="PDB ID for 3D Visualization"
+            )
+            
+            # Predict button
+            predict_btn = gr.Button("Predict Binding Sites")
+        
+        with gr.Column():
+            # Binding site predictions output
+            predictions_output = gr.Textbox(
+                label="Binding Site Predictions"
+            )
+            
+            # 3D Molecule visualization
+            molecule_output = Molecule3D(
+                label="Protein Structure", 
+                reps=reps
+            )
+    
+    # Prediction logic
+    predict_btn.click(
+        process_input, 
+        inputs=[sequence_input, pdb_input], 
+        outputs=[predictions_output, molecule_output]
+    )
+
+    # Add some example inputs
+    gr.Markdown("## Examples")
+    gr.Examples(
+        examples=[
+            ["MKVLWAALLVTFLAGCQAKVEQAVETEPEPELRQQTEWQSGQRWELALGRFWDYLRWVQTLSEQVQEELLSSQVTQELRALMDETMKELKAYKSELEEQLTPVAEETRARLSKELQAAQARLGADMEDVCGRLVQYRGEVQAMLGQSTEELRVRLASHLRKLRKRLLRDADDLQKRLAVYQAGAREGAERGLSAIRERLGPLVEQGRVRAATVGSLAGQPLQERAQAWGERLRARMEEMGSRTRDRLDEVKEQVAEVRAKLEEQAQQRL", "1ABC"],
+        ],
+        inputs=[sequence_input, pdb_input],
+        outputs=[predictions_output, molecule_output]
+    )
+
+demo.launch()

.ipynb_checkpoints/requirements-checkpoint.txt

@@ -9,4 +9,4 @@ scikit-learn>=0.24.0
 sentencepiece
 huggingface_hub>=0.15.0
 requests
-py3Dmol
+gradio_molecule3d

app.py

@@ -37,7 +37,7 @@ from scipy.special import expit
 
 import requests
 
-import py3Dmol
+from gradio_molecule3d import Molecule3D
 
 #import peft
 #from peft import get_peft_config, PeftModel, PeftConfig, inject_adapter_in_model, LoraConfig
@@ -46,6 +46,21 @@ import py3Dmol
 checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
 max_length = 1500
 
+# Default representations for molecule rendering
+reps = [
+    {
+        "model": 0,
+        "chain": "",
+        "resname": "",
+        "style": "cartoon",
+        "color": "spectrum",
+        "residue_range": "",
+        "around": 0,
+        "byres": False,
+        "visible": True
+    }
+]
+
 # Load model and move to device
 model, tokenizer = load_model(checkpoint, max_length)
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -53,7 +68,6 @@ model.to(device)
 model.eval()
 
 def create_dataset(tokenizer, seqs, labels, checkpoint):
-    
     tokenized = tokenizer(seqs, max_length=max_length, padding=False, truncation=True)
     dataset = Dataset.from_dict(tokenized)
     
@@ -68,7 +82,6 @@ def create_dataset(tokenizer, seqs, labels, checkpoint):
     return dataset
 
 def convert_predictions(input_logits):
-    
     all_probs = []
     for logits in input_logits:
         logits = logits.reshape(-1, 2)
@@ -78,13 +91,11 @@ def convert_predictions(input_logits):
     return np.concatenate(all_probs)
 
 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 predict_protein_sequence(test_one_letter_sequence):    
-    
     # Sanitize input sequence
     test_one_letter_sequence = test_one_letter_sequence.replace("O", "X") \
         .replace("B", "X").replace("U", "X") \
@@ -135,58 +146,88 @@ def predict_protein_sequence(test_one_letter_sequence):
     
     return result_str
 
-def fetch_and_display_pdb(pdb_id):
-    
+def fetch_pdb(pdb_id):
     try:
+        # Create a directory to store PDB files if it doesn't exist
+        os.makedirs('pdb_files', exist_ok=True)
+        
         # Fetch the PDB structure from RCSB
         pdb_url = f'https://files.rcsb.org/download/{pdb_id}.pdb'
-        response = requests.get(pdb_url)
-        
-        if response.status_code != 200:
-            return "Failed to load PDB structure. Please check the PDB ID."
+        pdb_path = f'pdb_files/{pdb_id}.pdb'
         
-        pdb_structure = response.text
+        # Download the file
+        response = requests.get(pdb_url)
         
-        # Prepare the 3D molecular visualization
-        visualization = f"""
-        <div id="container" style="width: 100%; height: 400px; position: relative;"></div>
-        <script src="https://3dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
-        <script>
-            let viewer = $3Dmol.createViewer(document.getElementById("container"));
-            viewer.addModel(`{pdb_structure}`, "pdb");
-            viewer.setStyle({{}}, {{"cartoon": {{"color": "spectrum"}}}});
-            viewer.zoomTo();
-            viewer.render();
-        </script>
-        """
-        return visualization
+        if response.status_code == 200:
+            with open(pdb_path, 'wb') as f:
+                f.write(response.content)
+            return pdb_path
+        else:
+            return None
     
     except Exception as e:
-        return f"Error visualizing PDB: {str(e)}"
+        print(f"Error fetching PDB: {e}")
+        return None
 
-def gradio_interface(sequence, pdb_id):
-    
+def process_input(sequence, pdb_id):
     # Predict binding sites
     binding_site_predictions = predict_protein_sequence(sequence)
     
-    # Fetch and visualize PDB structure
-    pdb_structure_html = fetch_and_display_pdb(pdb_id)
+    # Fetch PDB file
+    pdb_path = fetch_pdb(pdb_id)
     
-    return binding_site_predictions, pdb_structure_html
+    return binding_site_predictions, pdb_path
 
 # Create Gradio interface
-interface = gr.Interface(
-    fn=gradio_interface,
-    inputs=[
-        gr.Textbox(lines=2, placeholder="Enter protein sequence here...", label="Protein Sequence"),
-        gr.Textbox(lines=1, placeholder="Enter PDB ID here...", label="PDB ID for 3D Visualization")
-    ],
-    outputs=[
-        gr.Textbox(label="Binding Site Predictions"),
-        gr.HTML(label="3D Molecular Viewer")
-    ],
-    title="Protein Binding Site Prediction and 3D Structure Viewer",
-    description="Input a protein sequence to predict binding sites and view the protein structure in 3D using its PDB ID.",
-)
-
-interface.launch()
+with gr.Blocks() as demo:
+    gr.Markdown("# Protein Binding Site Prediction")
+    
+    with gr.Row():
+        with gr.Column():
+            # Sequence input
+            sequence_input = gr.Textbox(
+                lines=2, 
+                placeholder="Enter protein sequence here...", 
+                label="Protein Sequence"
+            )
+            
+            # PDB ID input
+            pdb_input = gr.Textbox(
+                lines=1, 
+                placeholder="Enter PDB ID here...", 
+                label="PDB ID for 3D Visualization"
+            )
+            
+            # Predict button
+            predict_btn = gr.Button("Predict Binding Sites")
+        
+        with gr.Column():
+            # Binding site predictions output
+            predictions_output = gr.Textbox(
+                label="Binding Site Predictions"
+            )
+            
+            # 3D Molecule visualization
+            molecule_output = Molecule3D(
+                label="Protein Structure", 
+                reps=reps
+            )
+    
+    # Prediction logic
+    predict_btn.click(
+        process_input, 
+        inputs=[sequence_input, pdb_input], 
+        outputs=[predictions_output, molecule_output]
+    )
+
+    # Add some example inputs
+    gr.Markdown("## Examples")
+    gr.Examples(
+        examples=[
+            ["MKVLWAALLVTFLAGCQAKVEQAVETEPEPELRQQTEWQSGQRWELALGRFWDYLRWVQTLSEQVQEELLSSQVTQELRALMDETMKELKAYKSELEEQLTPVAEETRARLSKELQAAQARLGADMEDVCGRLVQYRGEVQAMLGQSTEELRVRLASHLRKLRKRLLRDADDLQKRLAVYQAGAREGAERGLSAIRERLGPLVEQGRVRAATVGSLAGQPLQERAQAWGERLRARMEEMGSRTRDRLDEVKEQVAEVRAKLEEQAQQRL", "1ABC"],
+        ],
+        inputs=[sequence_input, pdb_input],
+        outputs=[predictions_output, molecule_output]
+    )
+
+demo.launch()

requirements.txt

@@ -9,4 +9,4 @@ scikit-learn>=0.24.0
 sentencepiece
 huggingface_hub>=0.15.0
 requests
-py3Dmol
+gradio_molecule3d