Update app.py

app.py

@@ -7,10 +7,11 @@ import torch
 import itertools
 from torch.utils.data import DataLoader
 from transformers import AutoTokenizer
+import shap
 
 sys.path.append("scripts/")
 from foldseek_util import get_struc_seq
-from utils import seed_everything
+from utils import seed_everything, save_pickle
 from models import PLTNUM_PreTrainedModel
 from datasets_ import PLTNUMDataset
 
@@ -26,6 +27,7 @@ class Config:
         self.task = "classification"
         self.sequence_col = "sequence"
         self.seed = 42
+        self.max_evals = 10
 
 
 
@@ -142,6 +144,71 @@ def predict(cfg, sequences):
     return predictions, [1 if x > 0.5 else 0 for x in predictions]
 
 
+
+def calculate_shap_values_with_pdb(model_choice, organism_choice, pdb_files, cfg=Config()):
+    input_sequences = []
+
+    for pdb_file in pdb_files:
+        pdb_path = pdb_file.name
+        os.system("chmod 777 bin/foldseek")
+        sequences = get_foldseek_seq(pdb_path)
+        sequence = sequences[2] if model_choice == "SaProt" else sequences[0]
+        input_sequences.append(sequence)
+
+    shap_values = calculate_shap_values_core(model_choice, organism_choice, input_sequences, cfg)
+
+    output_path = "/tmp/shap_values.pkl"
+    save_pickle(
+        output_path, shap_values
+    )
+
+    return output_path
+
+
+def calculate_shap_fn(texts, model, cfg):
+    if len(texts) == 1:
+        texts = texts[0]
+    else:
+        texts = texts.tolist()
+
+    inputs = cfg.tokenizer(
+        texts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=cfg.max_length,
+    )
+    inputs = {k: v.to(cfg.device) for k, v in inputs.items()}
+    with torch.no_grad():
+        outputs = model(inputs)
+        outputs = torch.sigmoid(outputs).detach().cpu().numpy()
+    return outputs
+
+
+def calculate_shap_values_core(model_choice, organism_choice, sequences, cfg=Config()):
+    cfg.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    seed_everything(cfg.seed)
+    tokenizer = AutoTokenizer.from_pretrained(
+        cfg.model_path, padding_side=cfg.padding_side
+    )
+    cfg.tokenizer = tokenizer
+
+    model = PLTNUM_PreTrainedModel.from_pretrained(cfg.model_path, cfg=cfg).to(cfg.device)
+    model.eval()
+
+    # build an explainer using a token masker
+    explainer = shap.Explainer(lambda x: calculate_shap_fn(x, model, cfg), cfg.tokenizer)
+
+    shap_values = explainer(
+        sequences,
+        batch_size=cfg.batch_size,
+        max_evals=cfg.max_evals,
+    )
+
+    return shap_values
+
+
 # Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -184,6 +251,17 @@ with gr.Blocks() as demo:
                 outputs=prediction_output,
             )
 
+            calculate_shap_values_button = gr.Button("Calculate SHAP Values")
+            shap_values_output = gr.File(
+                label="Download SHAP Values"
+            )
+            calculate_shap_values_button.click(
+                fn=calculate_shap_values_with_pdb,
+                inputs=[model_choice, organism_choice, pdb_files],
+                outputs=shap_values_output,
+            )
+
+
         with gr.TabItem("Enter Protein Sequence"):
             gr.Markdown("### Enter the protein sequence:")
             sequence = gr.Textbox(