Working loading and predicting with pretrained model + Added scaling in "pure" pytorch during forwarding

protac_degradation_predictor/__init__.py

@@ -16,6 +16,10 @@ from .optuna_utils import (
     hyperparameter_tuning_and_training,
     hyperparameter_tuning_and_training_sklearn,
 )
+from .protac_degradation_predictor import (
+    get_protac_active_proba,
+    is_protac_active,
+)
 
 __version__ = "0.0.1"
 __author__ = "Stefano Ribes"

protac_degradation_predictor/models/best_model_n0_random-epoch=6-val_acc=0.74-val_roc_auc=0.796.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52e12060ef0d21f4eb4d84570c708e8c8502a1fbe6ebcbae1d86044e23b77708
+size 101362856

protac_degradation_predictor/protac_degradation_predictor.py

@@ -1,5 +1,6 @@
 import pkg_resources
 import logging
+from typing import List
 
 from .pytorch_models import PROTAC_Model, load_model
 from .data_utils import (
@@ -15,37 +16,48 @@ from torch import sigmoid
 
 
 def get_protac_active_proba(
-        protac_smiles: str,
-        e3_ligase: str,
-        target_uniprot: str,
-        cell_line: str,
+        protac_smiles: str | List[str],
+        e3_ligase: str | List[str],
+        target_uniprot: str | List[str],
+        cell_line: str | List[str],
         device: str = 'cpu',
 ) -> bool:
-    ckpt_path = pkg_resources.resource_stream(__name__, 'data/model.ckpt')
+    
+    model_filename = 'best_model_n0_random-epoch=6-val_acc=0.74-val_roc_auc=0.796.ckpt'
+    ckpt_path = pkg_resources.resource_stream(__name__, f'models/{model_filename}')
     model = load_model(ckpt_path).to(device)
     protein2embedding = load_protein2embedding()
-    cell2embedding = load_cell2embedding()
+    cell2embedding = load_cell2embedding('data/cell2embedding.pkl')
 
     # Setup default embeddings
-    if e3_ligase not in config.e3_ligase2uniprot:
-        available_e3_ligases = ', '.join(list(config.e3_ligase2uniprot.keys()))
-        logging.warning(f"The E3 ligase {e3_ligase} is not in the database. Using the default E3 ligase. Available E3 ligases are: {available_e3_ligases}")
-    if target_uniprot not in protein2embedding:
-        logging.warning(f"The target protein {target_uniprot} is not in the database. Using the default target protein.")
-    if cell_line not in load_cell2embedding():
-        logging.warning(f"The cell line {cell_line} is not in the database. Using the default cell line.")
-
     default_protein_emb = np.zeros(config.protein_embedding_size)
     default_cell_emb = np.zeros(config.cell_embedding_size)
-    
+
     # Convert the E3 ligase to Uniprot ID
-    e3_ligase_uniprot = config.e3_ligase2uniprot.get(e3_ligase, '')
+    if isinstance(e3_ligase, list):
+        e3_ligase_uniprot = [config.e3_ligase2uniprot.get(e3, '') for e3 in e3_ligase]
+    else:
+        e3_ligase_uniprot = config.e3_ligase2uniprot.get(e3_ligase, '')
 
     # Get the embeddings
-    poi_emb = protein2embedding.get(target_uniprot, default_protein_emb)
-    e3_emb = protein2embedding.get(e3_ligase_uniprot, default_protein_emb)
-    cell_emb = cell2embedding.get(cell_line, default_cell_emb)
-    smiles_emb = get_fingerprint(protac_smiles)
+    if isinstance(protac_smiles, list):
+        # TODO: Add warning on missing entries?
+        poi_emb = [protein2embedding.get(t, default_protein_emb) for t in target_uniprot]
+        e3_emb = [protein2embedding.get(e3, default_protein_emb) for e3 in e3_ligase_uniprot]
+        cell_emb = [cell2embedding.get(cell_line, default_cell_emb) for cell_line in cell_line]
+        smiles_emb = [get_fingerprint(protac_smiles) for protac_smiles in protac_smiles]
+    else:
+        if e3_ligase not in config.e3_ligase2uniprot:
+            available_e3_ligases = ', '.join(list(config.e3_ligase2uniprot.keys()))
+            logging.warning(f"The E3 ligase {e3_ligase} is not in the database. Using the default E3 ligase. Available E3 ligases are: {available_e3_ligases}")
+        if target_uniprot not in protein2embedding:
+            logging.warning(f"The target protein {target_uniprot} is not in the database. Using the default target protein.")
+        if cell_line not in cell2embedding:
+            logging.warning(f"The cell line {cell_line} is not in the database. Using the default cell line.")
+        poi_emb = [protein2embedding.get(target_uniprot, default_protein_emb)]
+        e3_emb = [protein2embedding.get(e3_ligase_uniprot, default_protein_emb)]
+        cell_emb = [cell2embedding.get(cell_line, default_cell_emb)]
+        smiles_emb = [get_fingerprint(protac_smiles)]
 
     # Convert to torch tensors
     poi_emb = torch.tensor(poi_emb).to(device)
@@ -53,7 +65,18 @@ def get_protac_active_proba(
     cell_emb = torch.tensor(cell_emb).to(device)
     smiles_emb = torch.tensor(smiles_emb).to(device)
 
-    return model(poi_emb, e3_emb, cell_emb, smiles_emb).item()
+    pred = model(
+        poi_emb,
+        e3_emb,
+        cell_emb,
+        smiles_emb,
+        prescaled_embeddings=False, # Trigger automatic scaling
+    )
+
+    if isinstance(protac_smiles, list):
+        return sigmoid(pred).detach().numpy().flatten()
+    else:
+        return sigmoid(pred).item()
 
 
 def is_protac_active(
@@ -84,4 +107,4 @@ def is_protac_active(
         cell_line,
         device,
     )
-    return sigmoid(pred) > proba_threshold
+    return pred > proba_threshold

protac_degradation_predictor/pytorch_models.py

@@ -23,6 +23,7 @@ from torchmetrics import (
     MetricCollection,
 )
 from imblearn.over_sampling import SMOTE
+from sklearn.preprocessing import StandardScaler
 
 
 class PROTAC_Predictor(nn.Module):
@@ -239,8 +240,46 @@ class PROTAC_Model(pl.LightningModule):
             self.val_dataset.apply_scaling(self.scalers, use_single_scaler)
         if self.test_dataset:
             self.test_dataset.apply_scaling(self.scalers, use_single_scaler)
+    
+    def scale_tensor(
+            self,
+            tensor: torch.Tensor,
+            scaler: StandardScaler,
+    ) -> torch.Tensor:
+        """Scale a tensor using a scaler. This is done to avoid using numpy
+        arrays (and stay on the same device).
+        
+        Args:
+            tensor (torch.Tensor): The tensor to scale.
+            scaler (StandardScaler): The scaler to use.
 
-    def forward(self, poi_emb, e3_emb, cell_emb, smiles_emb):
+        Returns:
+            torch.Tensor: The scaled tensor.
+        """
+        tensor = tensor.float()
+        if scaler.with_mean:
+            tensor -= torch.tensor(scaler.mean_, dtype=tensor.dtype, device=tensor.device)
+        if scaler.with_std:
+            tensor /= torch.tensor(scaler.scale_, dtype=tensor.dtype, device=tensor.device)
+        return tensor
+
+    def forward(self, poi_emb, e3_emb, cell_emb, smiles_emb, prescaled_embeddings=True):
+        if not prescaled_embeddings:
+            if self.apply_scaling:
+                if self.join_embeddings == 'beginning':
+                    embeddings = self.scale_tensor(
+                        torch.hstack([smiles_emb, poi_emb, e3_emb, cell_emb]),
+                        self.scalers,
+                    )
+                    smiles_emb = embeddings[:, :self.smiles_emb_dim]
+                    poi_emb = embeddings[:, self.smiles_emb_dim:self.smiles_emb_dim+self.poi_emb_dim]
+                    e3_emb = embeddings[:, self.smiles_emb_dim+self.poi_emb_dim:self.smiles_emb_dim+2*self.poi_emb_dim]
+                    cell_emb = embeddings[:, -self.cell_emb_dim:]
+                else:
+                    poi_emb = self.scale_tensor(poi_emb, self.scalers['Uniprot'])
+                    e3_emb = self.scale_tensor(e3_emb, self.scalers['E3 Ligase Uniprot'])
+                    cell_emb = self.scale_tensor(cell_emb, self.scalers['Cell Line Identifier'])
+                    smiles_emb = self.scale_tensor(smiles_emb, self.scalers['Smiles'])
         return self.model(poi_emb, e3_emb, cell_emb, smiles_emb)
 
     def step(self, batch, batch_idx, stage):
@@ -451,7 +490,7 @@ def train_model(
         ),
         pl.callbacks.EarlyStopping(
             monitor='val_loss',
-            patience=5,
+            patience=10, # Original: 5
             mode='min',
             verbose=False,
         ),

setup.py

@@ -17,5 +17,5 @@ setuptools.setup(
         "Operating System :: OS Independent",
     ],
     include_package_data=True,
-    package_data={"": ["data/*.h5", "data/*.pkl", "data/*.csv"]},
+    package_data={"": ["data/*.h5", "data/*.pkl", "data/*.csv", "models/*.ckpt"]},
 )

tests/test_degradation_prediction.py

@@ -0,0 +1,45 @@
+import pytest
+import os
+import sys
+import logging
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from protac_degradation_predictor import (
+    get_protac_active_proba,
+    is_protac_active,
+)
+
+import torch
+
+
+def test_active_proba():
+    protac_smiles = 'Cc1ncsc1-c1ccc([C@H](C)NC(=O)[C@@H]2C[C@@H](O)CN2C(=O)[C@@H](NC(=O)CC(=O)N2CCN(CC[C@H](CSc3ccccc3)Nc3ccc(S(=O)(=O)NC(=O)c4ccc(N5CCN(CC6=C(c7ccc(Cl)cc7)CCC(C)(C)C6)CC5)cc4)cc3S(=O)(=O)C(F)(F)F)CC2)C(C)(C)C)cc1'
+    e3_ligase = 'VHL'
+    target_uniprot = 'Q07817'
+    cell_line = 'MOLT-4'
+    device = 'cpu'
+
+    active_prob = get_protac_active_proba(
+        protac_smiles=protac_smiles,
+        e3_ligase=e3_ligase,
+        target_uniprot=target_uniprot,
+        cell_line=cell_line,
+        device=device,
+    )
+
+    print(f'Active probability: {active_prob} (CPU)')
+
+    active_prob = get_protac_active_proba(
+        protac_smiles=[protac_smiles] * 16,
+        e3_ligase=[e3_ligase] * 16,
+        target_uniprot=[target_uniprot] * 16,
+        cell_line=[cell_line] * 16,
+        device='gpu' if torch.cuda.is_available() else 'cpu',
+    )
+
+    print(f'Active probability: {active_prob} (GPU)')
+
+
+def test_is_protac_active():
+    pass

tests/test_pytorch_model.py

@@ -40,19 +40,21 @@ def test_protac_predictor():
 def test_load_model(caplog):
     caplog.set_level(logging.WARNING)
 
+    model_filename = 'data/best_model_n0_random-epoch=6-val_acc=0.74-val_roc_auc=0.796.ckpt'
+
     model = PROTAC_Model.load_from_checkpoint(
-        'data/test_model.ckpt',
+        model_filename,
         map_location=torch.device("cpu") if not torch.cuda.is_available() else None,
     )
     # apply_scaling: true
     # batch_size: 8
     # cell_emb_dim: 768
     # disabled_embeddings: []
-    # dropout: 0.1498104322091649
+    # dropout: 0.11257777663560328
     # e3_emb_dim: 1024
     # hidden_dim: 768
     # join_embeddings: concat
-    # learning_rate: 4.881387978425994e-05
+    # learning_rate: 1.843233973932415e-05
     # poi_emb_dim: 1024
     # smiles_emb_dim: 224
     assert model.hidden_dim == 768
@@ -61,20 +63,25 @@ def test_load_model(caplog):
     assert model.e3_emb_dim == 1024
     assert model.cell_emb_dim == 768
     assert model.batch_size == 8
-    assert model.learning_rate == 4.881387978425994e-05
-    assert model.dropout == 0.1498104322091649
+    assert model.learning_rate == 1.843233973932415e-05
+    assert model.dropout == 0.11257777663560328
     assert model.join_embeddings == 'concat'
     assert model.disabled_embeddings == []
     assert model.apply_scaling == True
+    print(model.scalers)
 
 
 def test_checkpoint_file():
+    model_filename = 'data/best_model_n0_random-epoch=6-val_acc=0.74-val_roc_auc=0.796.ckpt'
     checkpoint = torch.load(
-        'data/test_model.ckpt',
+        model_filename,
         map_location=torch.device("cpu") if not torch.cuda.is_available() else None,
     )
     print(checkpoint.keys())
     print(checkpoint["hyper_parameters"])
     print([k for k, v in checkpoint["state_dict"].items()])
+    import pickle
+
+    print(pickle.loads(checkpoint['scalers']))
 
 pytest.main()