predict pipeline in dm and lm

Author: aditya0by0

chebai/models/base.py

@@ -232,7 +232,13 @@ def predict_step(
         Returns:
             Dict[str, Union[torch.Tensor, Any]]: The result of the prediction step.
         """
-        return self._execute(batch, batch_idx, self.test_metrics, prefix="", log=False)
+        assert isinstance(batch, XYData)
+        batch = batch.to(self.device)
+        data = self._process_batch(batch, batch_idx)
+        labels = data["labels"]
+        model_output = self(data, **data.get("model_kwargs", dict()))
+        pr, _ = self._get_prediction_and_labels(data, labels, model_output)
+        return pr
 
     def _execute(
         self,

chebai/preprocessing/datasets/base.py

@@ -339,8 +339,14 @@ def _load_data_from_file(self, path: str) -> List[Dict[str, Any]]:
             for d in tqdm.tqdm(self._load_dict(path), total=lines)
             if d["features"] is not None
         ]
+
+        return self._filter_to_token_limit(data)
+
+    def _filter_to_token_limit(
+        self, data: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
         # filter for missing features in resulting data, keep features length below token limit
-        data = [
+        return [
             val
             for val in data
             if val["features"] is not None
@@ -349,8 +355,6 @@ def _load_data_from_file(self, path: str) -> List[Dict[str, Any]]:
             )
         ]
 
-        return data
-
     def train_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]]:
         """
         Returns the train DataLoader.
@@ -400,10 +404,13 @@ def test_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]
         Returns:
             Union[DataLoader, List[DataLoader]]: A DataLoader object for test data.
         """
+
         return self.dataloader("test", shuffle=False, **kwargs)
 
     def predict_dataloader(
-        self, *args, **kwargs
+        self,
+        smiles_list: List[str],
+        **kwargs,
     ) -> Union[DataLoader, List[DataLoader]]:
         """
         Returns the predict DataLoader.
@@ -415,7 +422,21 @@ def predict_dataloader(
         Returns:
             Union[DataLoader, List[DataLoader]]: A DataLoader object for prediction data.
         """
-        return self.dataloader(self.prediction_kind, shuffle=False, **kwargs)
+
+        data = [
+            self.reader.to_data(
+                {"id": f"smiles_{idx}", "features": smiles, "labels": None}
+            )
+            for idx, smiles in enumerate(smiles_list)
+        ]
+        data = self._filter_to_token_limit(data)
+
+        return DataLoader(
+            data,
+            collate_fn=self.reader.collator,
+            batch_size=self.batch_size,
+            **kwargs,
+        )
 
     def prepare_data(self, *args, **kwargs) -> None:
         if self._prepare_data_flag != 1:

chebai/trainer/CustomTrainer.py

@@ -4,7 +4,7 @@
 
 import pandas as pd
 import torch
-from lightning import LightningModule, Trainer
+from lightning import Trainer
 from lightning.fabric.utilities.data import _set_sampler_epoch
 from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch.cli import instantiate_module
@@ -87,6 +87,7 @@ def predict_from_file(
         input_path: _PATH,
         save_to: _PATH = "predictions.csv",
         classes_path: Optional[_PATH] = None,
+        **kwargs,
     ) -> None:
         """
         Loads a model from a checkpoint and makes predictions on input data from a file.
@@ -114,6 +115,7 @@ def _predict_smiles(
         smiles: List[str],
         classes_path: Optional[_PATH] = None,
         save_to: _PATH = "predictions.csv",
+        **kwargs,
     ) -> torch.Tensor:
         """
         Predicts the output for a list of SMILES strings using the model.