predictors dont need targets fp

Author: aditya0by0

chebifier/prediction_models/electra_predictor.py

@@ -37,13 +37,13 @@ class ElectraPredictor(NNPredictor):
     def __init__(self, model_name: str, ckpt_path: str, **kwargs):
         super().__init__(model_name, ckpt_path, **kwargs)
         print(
-            f"Initialised Electra model {self.model_name} (device: {self._predictor.device})"
+            f"Initialised Electra model {self.model_name} (device: {self.predictor.device})"
         )
 
     def explain_smiles(self, smiles) -> dict:
         from chebai.preprocessing.reader import EMBEDDING_OFFSET
 
-        token_dict = self._predictor._dm.reader.to_data(
+        token_dict = self.predictor._dm.reader.to_data(
             dict(features=smiles, labels=None)
         )
         tokens = np.array(token_dict["features"]).astype(int).tolist()

chebifier/prediction_models/nn_predictor.py

@@ -16,38 +16,27 @@ def __init__(
         self,
         model_name: str,
         ckpt_path: str,
-        target_labels_path: str,
         **kwargs,
     ):
         super().__init__(model_name, **kwargs)
         self.batch_size = kwargs.get("batch_size", None)
         # If batch_size is not provided, it will be set to default batch size used during training in Predictor
-        self._predictor: Predictor = Predictor(ckpt_path, self.batch_size)
-        self.target_labels = [
-            line.strip() for line in open(target_labels_path, encoding="utf-8")
-        ]
-
-        # Sanity check - ensure that the number of classes predicted by the model matches the number of target labels
-        # TODO: In future, we can include the target labels in the model metadata and avoid this.
-        raw_preds = self._predictor.predict_smiles(["CO"])
-        assert len(raw_preds[0]) == len(
-            self.target_labels
-        ), "Number of predicted classes does not match number of target labels."
+        self.predictor: Predictor = Predictor(ckpt_path, self.batch_size)
 
     @modelwise_smiles_lru_cache.batch_decorator
     def predict_smiles_list(self, smiles_list: list[str]) -> list:
         """
         Returns a list with the length of smiles_list, each element is
         either None (=failure) or a dictionary of classes and predicted values.
         """
-        raw_preds: Tensor = self._predictor.predict_smiles(smiles_list)
+        raw_preds: Tensor = self.predictor.predict_smiles(smiles_list)
         if raw_preds is not None:
             preds = [
                 (
                     {
                         label: pred
                         for label, pred in zip(
-                            self.target_labels, raw_preds[i].tolist()
+                            self.predictor._classification_labels, raw_preds[i].tolist()
                         )
                     }
                 )
@@ -56,3 +45,10 @@ def predict_smiles_list(self, smiles_list: list[str]) -> list:
             return preds
         else:
             return [None for _ in smiles_list]
+
+    def calculate_results(self, batch):
+        collator = self.predictor._dm.reader.COLLATOR()
+        dat = self.predictor._model._process_batch(
+            collator(batch).to(self.predictor.device), 0
+        )
+        return self.predictor._model(dat, **dat["model_kwargs"])