handle augment electra and old ckpt files

Author: aditya0by0

chebai/result/prediction.py

@@ -37,6 +37,13 @@ def __init__(
 
         self._dm_hparams = ckpt_file["datamodule_hyper_parameters"]
         self._dm_hparams.pop("splits_file_path")
+        self._dm_hparams.pop("augment_smiles", None)
+        self._dm_hparams.pop("aug_smiles_variations", None)
+        assert "_class_path" in self._dm_hparams, (
+            "Datamodule hyperparameters must include a '_class_path' key.\n"
+            "Hence, either the checkpoint is corrupted or "
+            "it was not saved properly with latest lightning version"
+        )
         self._dm: XYBaseDataModule = instantiate_module(
             XYBaseDataModule, self._dm_hparams
         )