gather_predictions will return predicted_classes_dict

Author: aditya0by0

chebifier/ensemble/base_ensemble.py

@@ -1,8 +1,8 @@
 import os
 from abc import ABC
+
 import torch
 import tqdm
-from rdkit import Chem
 
 from chebifier.prediction_models.base_predictor import BasePredictor
 from chebifier.prediction_models.chemlog_predictor import ChemLogPredictor
@@ -12,11 +12,11 @@
 MODEL_TYPES = {
     "electra": ElectraPredictor,
     "resgated": ResGatedPredictor,
-    "chemlog": ChemLogPredictor
+    "chemlog": ChemLogPredictor,
 }
 
-class BaseEnsemble(ABC):
 
+class BaseEnsemble(ABC):
     def __init__(self, model_configs: dict):
         self.models = []
         self.positive_prediction_threshold = 0.5
@@ -37,22 +37,30 @@ def gather_predictions(self, smiles_list):
                 if logits_for_smiles is not None:
                     for cls in logits_for_smiles:
                         predicted_classes.add(cls)
-        print(f"Sorting predictions...")
+        print("Sorting predictions...")
         predicted_classes = sorted(list(predicted_classes))
         predicted_classes_dict = {cls: i for i, cls in enumerate(predicted_classes)}
-        ordered_logits = torch.zeros(len(smiles_list), len(predicted_classes), len(self.models)) * torch.nan
+        ordered_logits = (
+            torch.zeros(len(smiles_list), len(predicted_classes), len(self.models))
+            * torch.nan
+        )
         for i, model_prediction in enumerate(model_predictions):
-            for j, logits_for_smiles in tqdm.tqdm(enumerate(model_prediction),
-                                                 total=len(model_prediction),
-                                                 desc=f"Sorting predictions for {self.models[i].model_name}"):
+            for j, logits_for_smiles in tqdm.tqdm(
+                enumerate(model_prediction),
+                total=len(model_prediction),
+                desc=f"Sorting predictions for {self.models[i].model_name}",
+            ):
                 if logits_for_smiles is not None:
                     for cls in logits_for_smiles:
-                        ordered_logits[j, predicted_classes_dict[cls], i] = logits_for_smiles[cls]
+                        ordered_logits[j, predicted_classes_dict[cls], i] = (
+                            logits_for_smiles[cls]
+                        )
 
-        return ordered_logits, predicted_classes
+        return ordered_logits, predicted_classes_dict
 
-
-    def consolidate_predictions(self, predictions, predicted_classes, classwise_weights, **kwargs):
+    def consolidate_predictions(
+        self, predictions, predicted_classes, classwise_weights, **kwargs
+    ):
         """
         Aggregates predictions from multiple models using weighted majority voting.
         Optimized version using tensor operations instead of for loops.
@@ -74,7 +82,9 @@ def consolidate_predictions(self, predictions, predicted_classes, classwise_weig
         positive_mask = (predictions > 0.5) & valid_predictions
         negative_mask = (predictions < 0.5) & valid_predictions
 
-        confidence = 2 * torch.abs(predictions.nan_to_num() - self.positive_prediction_threshold)
+        confidence = 2 * torch.abs(
+            predictions.nan_to_num() - self.positive_prediction_threshold
+        )
 
         # Extract positive and negative weights
         pos_weights = classwise_weights[0]  # Shape: (num_classes, num_models)
@@ -83,26 +93,34 @@ def consolidate_predictions(self, predictions, predicted_classes, classwise_weig
         # Calculate weighted predictions using broadcasting
         # predictions shape: (num_smiles, num_classes, num_models)
         # weights shape: (num_classes, num_models)
-        positive_weighted = positive_mask.float() * confidence * pos_weights.unsqueeze(0)
-        negative_weighted = negative_mask.float() * confidence * neg_weights.unsqueeze(0)
+        positive_weighted = (
+            positive_mask.float() * confidence * pos_weights.unsqueeze(0)
+        )
+        negative_weighted = (
+            negative_mask.float() * confidence * neg_weights.unsqueeze(0)
+        )
 
         # Sum over models dimension
         positive_sum = positive_weighted.sum(dim=2)  # Shape: (num_smiles, num_classes)
         negative_sum = negative_weighted.sum(dim=2)  # Shape: (num_smiles, num_classes)
 
         # Determine which classes to include for each SMILES
         net_score = positive_sum - negative_sum  # Shape: (num_smiles, num_classes)
-        class_decisions = (net_score > 0) & has_valid_predictions  # Shape: (num_smiles, num_classes)
+        class_decisions = (
+            net_score > 0
+        ) & has_valid_predictions  # Shape: (num_smiles, num_classes)
 
         # Convert tensor decisions to result list using list comprehension for efficiency
         result = [
-            [class_indices[idx.item()] for idx in torch.nonzero(class_decisions[i], as_tuple=True)[0]]
+            [
+                class_indices[idx.item()]
+                for idx in torch.nonzero(class_decisions[i], as_tuple=True)[0]
+            ]
             for i in range(num_smiles)
         ]
 
         return result
 
-
     def calculate_classwise_weights(self, predicted_classes):
         """No weights, simple majority voting"""
         positive_weights = torch.ones(len(predicted_classes), len(self.models))
@@ -114,18 +132,26 @@ def predict_smiles_list(self, smiles_list, load_preds_if_possible=True) -> list:
         preds_file = f"predictions_by_model_{'_'.join(model.model_name for model in self.models)}.pt"
         predicted_classes_file = f"predicted_classes_{'_'.join(model.model_name for model in self.models)}.txt"
         if not load_preds_if_possible or not os.path.isfile(preds_file):
-            ordered_predictions, predicted_classes = self.gather_predictions(smiles_list)
+            ordered_predictions, predicted_classes = self.gather_predictions(
+                smiles_list
+            )
             # save predictions
             torch.save(ordered_predictions, preds_file)
             with open(predicted_classes_file, "w") as f:
                 for cls in predicted_classes:
                     f.write(f"{cls}\n")
         else:
-            print(f"Loading predictions from {preds_file} and label indexes from {predicted_classes_file}")
+            print(
+                f"Loading predictions from {preds_file} and label indexes from {predicted_classes_file}"
+            )
             ordered_predictions = torch.load(preds_file)
             with open(predicted_classes_file, "r") as f:
-                predicted_classes = {line.strip(): i for i, line in enumerate(f.readlines())}
+                predicted_classes = {
+                    line.strip(): i for i, line in enumerate(f.readlines())
+                }
 
         classwise_weights = self.calculate_classwise_weights(predicted_classes)
-        aggregated_predictions = self.consolidate_predictions(ordered_predictions, predicted_classes, classwise_weights)
+        aggregated_predictions = self.consolidate_predictions(
+            ordered_predictions, predicted_classes, classwise_weights
+        )
         return aggregated_predictions