Update

.ipynb_checkpoints/app-checkpoint.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 import gradio as gr
 import requests
-from Bio.PDB import PDBParser, MMCIFParser, PDBIO
+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
@@ -9,7 +9,6 @@ import numpy as np
 import os
 from gradio_molecule3d import Molecule3D
 
-
 from model_loader import load_model
 
 import torch
@@ -21,31 +20,13 @@ import re
 import pandas as pd
 import copy
 
-import transformers, datasets
-from transformers import AutoTokenizer
-from transformers import DataCollatorForTokenClassification
+import transformers
+from transformers import AutoTokenizer, DataCollatorForTokenClassification
 
 from datasets import Dataset
 
 from scipy.special import expit
 
-from datetime import datetime
-import gradio as gr
-import requests
-from Bio.PDB import PDBParser, MMCIFParser, PDBIO
-from Bio.PDB.Polypeptide import is_aa
-from Bio.SeqUtils import seq1
-from typing import Optional, Tuple
-import numpy as np
-import os
-from gradio_molecule3d import Molecule3D
-
-import re
-import pandas as pd
-import copy
-
-from scipy.special import expit
-
 
 # Load model and move to device
 checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
@@ -55,24 +36,6 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)
 model.eval()
 
-from datetime import datetime
-import gradio as gr
-import requests
-from Bio.PDB import PDBParser, MMCIFParser, PDBIO
-from Bio.PDB.Polypeptide import is_aa
-from Bio.SeqUtils import seq1
-from Bio.PDB import Select
-from typing import Optional, Tuple
-import numpy as np
-import os
-from gradio_molecule3d import Molecule3D
-
-import re
-import pandas as pd
-import copy
-
-from scipy.special import expit
-
 def normalize_scores(scores):
     min_score = np.min(scores)
     max_score = np.max(scores)
@@ -232,7 +195,13 @@ def process_pdb(pdb_id_or_file, segment):
     sequence = "".join(seq1(res.resname) for res in protein_residues)
     sequence_id = [res.id[1] for res in protein_residues]
     
-    scores = np.random.rand(len(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()
+    
+    # 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
@@ -404,8 +373,8 @@ def molecule(input_pdb, residue_scores=None, segment='A'):
                                     position: atom, 
                                     backgroundColor: 'mintcream', 
                                     fontColor: 'black',
-                                    fontSize: 12,
-                                    padding: 2
+                                    fontSize: 18,
+                                    padding: 4
                                 }}
                             );
                         }}
@@ -464,7 +433,7 @@ with gr.Blocks() as demo:
 
     molecule_output = gr.HTML(label="Protein Structure")
     explanation_vis = gr.Markdown("""
-    Residues with a score > 0.5 are considered binding sites and represented as sticks with the score dependent colorcoding:
+    Residues with a score > 0.5 are represented as sticks with a score dependent colorcoding:
     - 0.5-0.6: blue  
     - 0.6–0.7: light blue  
     - 0.7–0.8: white

app.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 import gradio as gr
 import requests
-from Bio.PDB import PDBParser, MMCIFParser, PDBIO
+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
@@ -9,7 +9,6 @@ import numpy as np
 import os
 from gradio_molecule3d import Molecule3D
 
-
 from model_loader import load_model
 
 import torch
@@ -21,31 +20,13 @@ import re
 import pandas as pd
 import copy
 
-import transformers, datasets
-from transformers import AutoTokenizer
-from transformers import DataCollatorForTokenClassification
+import transformers
+from transformers import AutoTokenizer, DataCollatorForTokenClassification
 
 from datasets import Dataset
 
 from scipy.special import expit
 
-from datetime import datetime
-import gradio as gr
-import requests
-from Bio.PDB import PDBParser, MMCIFParser, PDBIO
-from Bio.PDB.Polypeptide import is_aa
-from Bio.SeqUtils import seq1
-from typing import Optional, Tuple
-import numpy as np
-import os
-from gradio_molecule3d import Molecule3D
-
-import re
-import pandas as pd
-import copy
-
-from scipy.special import expit
-
 
 # Load model and move to device
 checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
@@ -55,24 +36,6 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)
 model.eval()
 
-from datetime import datetime
-import gradio as gr
-import requests
-from Bio.PDB import PDBParser, MMCIFParser, PDBIO
-from Bio.PDB.Polypeptide import is_aa
-from Bio.SeqUtils import seq1
-from Bio.PDB import Select
-from typing import Optional, Tuple
-import numpy as np
-import os
-from gradio_molecule3d import Molecule3D
-
-import re
-import pandas as pd
-import copy
-
-from scipy.special import expit
-
 def normalize_scores(scores):
     min_score = np.min(scores)
     max_score = np.max(scores)
@@ -410,8 +373,8 @@ def molecule(input_pdb, residue_scores=None, segment='A'):
                                     position: atom, 
                                     backgroundColor: 'mintcream', 
                                     fontColor: 'black',
-                                    fontSize: 12,
-                                    padding: 2
+                                    fontSize: 18,
+                                    padding: 4
                                 }}
                             );
                         }}
@@ -470,7 +433,7 @@ with gr.Blocks() as demo:
 
     molecule_output = gr.HTML(label="Protein Structure")
     explanation_vis = gr.Markdown("""
-    Residues with a score > 0.5 are considered binding sites and represented as sticks with the score dependent colorcoding:
+    Residues with a score > 0.5 are represented as sticks with a score dependent colorcoding:
     - 0.5-0.6: blue  
     - 0.6–0.7: light blue  
     - 0.7–0.8: white